tf.keras.config.set_x64_mode

Sets how Keras handles 64-bit dtypes.

tf.keras.config.set_x64_mode(
    mode
)

64-bit dtypes are int64, uint64, float64 and complex128.

mode One of:

  • "off" (default): the result dtype of every op is capped to 32-bit, e.g. adding two int64 tensors gives an int32 tensor.
  • "explicit_x64": a 64-bit input gives a 64-bit output. Inputs that are all 32-bit, Python scalars and default dtypes stay 32-bit. This mirrors JAX's jax_explicit_x64_dtypes="allow", which Keras turns on for you on the JAX backend.
  • "enable_x64": not supported yet.

The mode can also be set with the KERAS_X64_MODE environment variable or the "x64_mode" field of ~/.keras/keras.json. Set it before creating tensors or models. Supported on the jax, numpy and torch backends. On the JAX backend, Keras turns on the matching JAX flag. If JAX's own 64-bit flags are already on and don't match mode, Keras warns and keeps mode.

Example:

keras.config.set_x64_mode("explicit_x64")
x = keras.ops.ones((2,), dtype="int64")
keras.ops.add(x, x).dtype
'int64'

ValueError If mode is unknown, not supported yet, or not supported on the current backend.