def _sharding_constraint_impl(x, partitions):
  # TODO(skye): can we also prevent this from being called in other
  # non-sharded_jit contexts? (e.g. pmap, control flow)
  raise NotImplementedError(
      "with_sharding_constraint() should only be called inside sharded_jit()")