<p dir="auto">The following code:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="@functools.partial(jax.custom_vjp, nondiff_argnums=(1,))
def _precise_mean(x, axis):
  with jax.experimental.enable_x64():
    return jnp.mean(x.astype(&quot;float64&quot;), axis=axis).astype(x.dtype)


def _precise_mean_fwd(x, axis):
  return _precise_mean(x, axis), (x,)
  
def _precise_mean_bwd(axis, res, g):
  x, = res
  dims = tuple(i for i in range(x.ndim) if i not in axis)
  factor = 1
  for a in axis:
    factor *= x.shape[a]
  g = g / factor
  print(g.dtype)
  return lax.broadcast_in_dim(g.astype(x.dtype), x.shape, dims),

_precise_mean.defvjp(_precise_mean_fwd, _precise_mean_bwd)


def precise_mean(x, axis=None):
  if axis is None:
    axis = tuple(range(x.ndim))
  elif isinstance(axis, int):
    axis = (axis,)
  # Take care of negative values
  axis = tuple((x.ndim + a) % x.ndim for a in axis)
  if not all(isinstance(a, int) for a in axis):
    raise ValueError(&quot;All axis must be integers.&quot;)
  return _precise_mean(x, axis)

@jax.jit
def f(x, a, b):
  x = jnp.matmul(x, a) + b[None]
  y = jnp.sum(x, axis=-1)
  return precise_mean(y, axis=0)
  
x = jnp.ones([5000, 100])
a = jnp.ones([100, 100])
b = jnp.zeros([100])
print(f(x, a, b).dtype)"><pre class="notranslate"><code class="notranslate">@functools.partial(jax.custom_vjp, nondiff_argnums=(1,))
def _precise_mean(x, axis):
  with jax.experimental.enable_x64():
    return jnp.mean(x.astype("float64"), axis=axis).astype(x.dtype)


def _precise_mean_fwd(x, axis):
  return _precise_mean(x, axis), (x,)
  
def _precise_mean_bwd(axis, res, g):
  x, = res
  dims = tuple(i for i in range(x.ndim) if i not in axis)
  factor = 1
  for a in axis:
    factor *= x.shape[a]
  g = g / factor
  print(g.dtype)
  return lax.broadcast_in_dim(g.astype(x.dtype), x.shape, dims),

_precise_mean.defvjp(_precise_mean_fwd, _precise_mean_bwd)


def precise_mean(x, axis=None):
  if axis is None:
    axis = tuple(range(x.ndim))
  elif isinstance(axis, int):
    axis = (axis,)
  # Take care of negative values
  axis = tuple((x.ndim + a) % x.ndim for a in axis)
  if not all(isinstance(a, int) for a in axis):
    raise ValueError("All axis must be integers.")
  return _precise_mean(x, axis)

@jax.jit
def f(x, a, b):
  x = jnp.matmul(x, a) + b[None]
  y = jnp.sum(x, axis=-1)
  return precise_mean(y, axis=0)
  
x = jnp.ones([5000, 100])
a = jnp.ones([100, 100])
b = jnp.zeros([100])
print(f(x, a, b).dtype)
</code></pre></div>
<p dir="auto">Throws out the following error:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="UnfilteredStackTrace: RuntimeError: Internal: Seen floating point types of different precisions in %reduce.23 = f32[] reduce(f64[5000]{0} %convert.17, f32[] %constant.18), dimensions={0}, to_apply=%primitive_computation_add.19, metadata={op_type=&quot;reduce_sum&quot; op_name=&quot;jit(f)/custom_vjp_call_jaxpr/reduce_sum[ axes=(0,) ]&quot; source_file=&quot;&lt;ipython-input-219-aea426c526ae&gt;&quot; source_line=13}, but mixed precision is disallowed.

Failed after pipeline-start

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified."><pre class="notranslate"><code class="notranslate">UnfilteredStackTrace: RuntimeError: Internal: Seen floating point types of different precisions in %reduce.23 = f32[] reduce(f64[5000]{0} %convert.17, f32[] %constant.18), dimensions={0}, to_apply=%primitive_computation_add.19, metadata={op_type="reduce_sum" op_name="jit(f)/custom_vjp_call_jaxpr/reduce_sum[ axes=(0,) ]" source_file="&lt;ipython-input-219-aea426c526ae&gt;" source_line=13}, but mixed precision is disallowed.

Failed after pipeline-start

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.
</code></pre></div>
<p dir="auto">It is unclear at the moment how to use the <code class="notranslate">enable_x64</code> contest when using <code class="notranslate">jax.jit</code>.</p>
<p dir="auto">minimal repro:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from jax import jit, experimental

@jit
def f(a):
  with experimental.enable_x64():
    return 1 + a
f(1)"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">jit</span>, <span class="pl-s1">experimental</span>

<span class="pl-en">@<span class="pl-s1">jit</span></span>
<span class="pl-k">def</span> <span class="pl-en">f</span>(<span class="pl-s1">a</span>):
  <span class="pl-k">with</span> <span class="pl-s1">experimental</span>.<span class="pl-en">enable_x64</span>():
    <span class="pl-k">return</span> <span class="pl-c1">1</span> <span class="pl-c1">+</span> <span class="pl-s1">a</span>
<span class="pl-en">f</span>(<span class="pl-c1">1</span>)</pre></div>
<div class="highlight highlight-text-python-traceback notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="RuntimeError: Invalid argument: Binary op add with different element types: s64[] and s32[]."><pre class="notranslate"><span class="pl-en">RuntimeError</span>: <span class="pl-s">Invalid argument: Binary op add with different element types: s64[] and s32[].</span></pre></div>