google-meridian 1.2.1__py3-none-any.whl → 1.3.1__py3-none-any.whl

meridian/backend/__init__.py

@@ -18,6 +18,7 @@ import abc
 import functools
 import os
 from typing import Any, Optional, Sequence, Tuple, TYPE_CHECKING, Union
+import warnings
 
 from meridian.backend import config
 import numpy as np
@@ -37,6 +38,7 @@ _DEFAULT_INT = "int64"
 _TENSORFLOW_TILE_KEYWORD = "multiples"
 _JAX_TILE_KEYWORD = "reps"
 
+_ARG_AUTOGRAPH = "autograph"
 _ARG_JIT_COMPILE = "jit_compile"
 _ARG_STATIC_ARGNUMS = "static_argnums"
 _ARG_STATIC_ARGNAMES = "static_argnames"
@@ -134,6 +136,54 @@ def _resolve_dtype(dtype: Optional[Any], *args: Any) -> str:
 
 
 # --- Private Backend-Specific Implementations ---
+def _jax_stabilize_rf_roi_grid(
+    spend_grid: np.ndarray,
+    outcome_grid: np.ndarray,
+    n_rf_channels: int,
+) -> np.ndarray:
+  """Stabilizes the RF ROI grid for JAX using a stable index lookup."""
+  new_outcome_grid = outcome_grid.copy()
+  rf_slice = slice(-n_rf_channels, None)
+  rf_spend_grid = spend_grid[:, rf_slice]
+  rf_outcome_grid = new_outcome_grid[:, rf_slice]
+
+  last_valid_indices = np.sum(~np.isnan(rf_spend_grid), axis=0) - 1
+  channel_indices = np.arange(n_rf_channels)
+
+  ref_spend = rf_spend_grid[last_valid_indices, channel_indices]
+  ref_outcome = rf_outcome_grid[last_valid_indices, channel_indices]
+
+  rf_roi = np.divide(
+      ref_outcome,
+      ref_spend,
+      out=np.zeros_like(ref_outcome, dtype=np.float64),
+      where=(ref_spend != 0),
+  )
+
+  new_outcome_grid[:, rf_slice] = rf_roi * rf_spend_grid
+  return new_outcome_grid
+
+
+def _tf_stabilize_rf_roi_grid(
+    spend_grid: np.ndarray,
+    outcome_grid: np.ndarray,
+    n_rf_channels: int,
+) -> np.ndarray:
+  """Stabilizes the RF ROI grid for TF using nanmax logic."""
+  new_outcome_grid = outcome_grid.copy()
+  rf_slice = slice(-n_rf_channels, None)
+  rf_outcome_max = np.nanmax(new_outcome_grid[:, rf_slice], axis=0)
+  rf_spend_max = np.nanmax(spend_grid[:, rf_slice], axis=0)
+
+  rf_roi = np.divide(
+      rf_outcome_max,
+      rf_spend_max,
+      out=np.zeros_like(rf_outcome_max, dtype=np.float64),
+      where=(rf_spend_max != 0),
+  )
+
+  new_outcome_grid[:, rf_slice] = rf_roi * spend_grid[:, rf_slice]
+  return new_outcome_grid
 
 
 def _jax_arange(
@@ -171,7 +221,7 @@ def _tf_arange(
 
 def _jax_cast(x: Any, dtype: Any) -> "_jax.Array":
   """JAX implementation for cast."""
-  return x.astype(dtype)
+  return jax_ops.asarray(x, dtype=dtype)
 
 
 def _jax_divide_no_nan(x, y):
@@ -205,6 +255,7 @@ def _jax_function_wrapper(func=None, **kwargs):
   jit_kwargs = kwargs.copy()
 
   jit_kwargs.pop(_ARG_JIT_COMPILE, None)
+  jit_kwargs.pop(_ARG_AUTOGRAPH, None)
 
   if _ARG_STATIC_ARGNUMS in jit_kwargs:
     if not jit_kwargs[_ARG_STATIC_ARGNUMS]:
@@ -255,17 +306,132 @@ def _jax_numpy_function(*args, **kwargs):  # pylint: disable=unused-argument
   )
 
 
-def _jax_make_tensor_proto(*args, **kwargs):  # pylint: disable=unused-argument
-  raise NotImplementedError(
-      "backend.make_tensor_proto is not implemented for the JAX backend."
-  )
+def _jax_make_tensor_proto(values, dtype=None, shape=None):  # pylint: disable=unused-argument
+  """JAX implementation for make_tensor_proto."""
+  # pylint: disable=g-direct-tensorflow-import
+  from tensorflow.core.framework import tensor_pb2
+  from tensorflow.core.framework import tensor_shape_pb2
+  from tensorflow.core.framework import types_pb2
+  # pylint: enable=g-direct-tensorflow-import
 
+  if not isinstance(values, np.ndarray):
+    values = np.array(values)
 
-def _jax_make_ndarray(*args, **kwargs):  # pylint: disable=unused-argument
-  raise NotImplementedError(
-      "backend.make_ndarray is not implemented for the JAX backend."
+  if dtype:
+    numpy_dtype = np.dtype(dtype)
+    values = values.astype(numpy_dtype)
+  else:
+    numpy_dtype = values.dtype
+
+  dtype_map = {
+      np.dtype(np.float16): types_pb2.DT_HALF,
+      np.dtype(np.float32): types_pb2.DT_FLOAT,
+      np.dtype(np.float64): types_pb2.DT_DOUBLE,
+      np.dtype(np.int32): types_pb2.DT_INT32,
+      np.dtype(np.uint8): types_pb2.DT_UINT8,
+      np.dtype(np.uint16): types_pb2.DT_UINT16,
+      np.dtype(np.uint32): types_pb2.DT_UINT32,
+      np.dtype(np.uint64): types_pb2.DT_UINT64,
+      np.dtype(np.int16): types_pb2.DT_INT16,
+      np.dtype(np.int8): types_pb2.DT_INT8,
+      np.dtype(np.int64): types_pb2.DT_INT64,
+      np.dtype(np.complex64): types_pb2.DT_COMPLEX64,
+      np.dtype(np.complex128): types_pb2.DT_COMPLEX128,
+      np.dtype(np.bool_): types_pb2.DT_BOOL,
+      # Note: String types are handled outside the map.
+  }
+  proto_dtype = dtype_map.get(numpy_dtype)
+  if proto_dtype is None and numpy_dtype.kind in ("S", "U"):
+    proto_dtype = types_pb2.DT_STRING
+
+  if proto_dtype is None:
+    raise TypeError(
+        f"Unsupported dtype for TensorProto conversion: {numpy_dtype}"
+    )
+
+  proto = tensor_pb2.TensorProto(
+      dtype=proto_dtype,
+      tensor_shape=tensor_shape_pb2.TensorShapeProto(
+          dim=[
+              tensor_shape_pb2.TensorShapeProto.Dim(size=d)
+              for d in values.shape
+          ]
+      ),
   )
 
+  proto.tensor_content = values.tobytes()
+  return proto
+
+
+def _jax_make_ndarray(proto):
+  """JAX implementation for make_ndarray."""
+  # pylint: disable=g-direct-tensorflow-import
+  from tensorflow.core.framework import types_pb2
+  # pylint: enable=g-direct-tensorflow-import
+
+  dtype_map = {
+      types_pb2.DT_HALF: np.float16,
+      types_pb2.DT_FLOAT: np.float32,
+      types_pb2.DT_DOUBLE: np.float64,
+      types_pb2.DT_INT32: np.int32,
+      types_pb2.DT_UINT8: np.uint8,
+      types_pb2.DT_UINT16: np.uint16,
+      types_pb2.DT_UINT32: np.uint32,
+      types_pb2.DT_UINT64: np.uint64,
+      types_pb2.DT_INT16: np.int16,
+      types_pb2.DT_INT8: np.int8,
+      types_pb2.DT_INT64: np.int64,
+      types_pb2.DT_COMPLEX64: np.complex64,
+      types_pb2.DT_COMPLEX128: np.complex128,
+      types_pb2.DT_BOOL: np.bool_,
+      types_pb2.DT_STRING: np.bytes_,
+  }
+  if proto.dtype not in dtype_map:
+    raise TypeError(f"Unsupported TensorProto dtype: {proto.dtype}")
+
+  shape = [d.size for d in proto.tensor_shape.dim]
+  dtype = dtype_map[proto.dtype]
+
+  if proto.tensor_content:
+    num_elements = np.prod(shape).item() if shape else 0
+    # When deserializing a string from tensor_content, the itemsize is not
+    # explicitly stored. We must infer it from the content length and shape.
+    if dtype == np.bytes_ and num_elements > 0:
+      content_len = len(proto.tensor_content)
+      itemsize = content_len // num_elements
+      if itemsize * num_elements != content_len:
+        raise ValueError(
+            "Tensor content size is not a multiple of the number of elements"
+            " for string dtype."
+        )
+      dtype = np.dtype(f"S{itemsize}")
+
+    return (
+        np.frombuffer(proto.tensor_content, dtype=dtype).copy().reshape(shape)
+    )
+
+  # Fallback for protos that store data in val fields instead of tensor_content.
+  if dtype == np.float32:
+    val_field = proto.float_val
+  elif dtype == np.float64:
+    val_field = proto.double_val
+  elif dtype == np.int32:
+    val_field = proto.int_val
+  elif dtype == np.int64:
+    val_field = proto.int64_val
+  elif dtype == np.bool_:
+    val_field = proto.bool_val
+  else:
+    if proto.string_val:
+      return np.array(proto.string_val, dtype=np.bytes_).reshape(shape)
+    if not any(shape):
+      return np.array([], dtype=dtype).reshape(shape)
+    raise TypeError(
+        f"Unsupported dtype for TensorProto value field fallback: {dtype}"
+    )
+
+  return np.array(val_field, dtype=dtype).reshape(shape)
+
 
 def _jax_get_indices_where(condition):
   """JAX implementation for get_indices_where."""
@@ -324,6 +490,7 @@ def _jax_boolean_mask(tensor, mask, axis=None):
 
   if axis is None:
     axis = 0
+  mask = jnp.asarray(mask)
   tensor_swapped = jnp.moveaxis(tensor, axis, 0)
   masked = tensor_swapped[mask]
   return jnp.moveaxis(masked, 0, axis)
@@ -336,17 +503,26 @@ def _tf_boolean_mask(tensor, mask, axis=None):
   return tf.boolean_mask(tensor, mask, axis=axis)
 
 
-def _jax_gather(params, indices):
-  """JAX implementation for gather."""
-  # JAX uses standard array indexing for gather operations.
-  return params[indices]
+def _jax_gather(params, indices, axis=0):
+  """JAX implementation for gather with axis support."""
+  import jax.numpy as jnp
+
+  if isinstance(params, (list, tuple)):
+    params = np.array(params)
+
+  # JAX can't JIT-compile operations on string or object arrays. We detect
+  # these types and fall back to standard NumPy operations.
+  if isinstance(params, np.ndarray) and params.dtype.kind in ("S", "U", "O"):
+    return np.take(params, np.asarray(indices), axis=axis)
+
+  return jnp.take(params, indices, axis=axis)
 
 
-def _tf_gather(params, indices):
+def _tf_gather(params, indices, axis=0):
   """TensorFlow implementation for gather."""
   import tensorflow as tf
 
-  return tf.gather(params, indices)
+  return tf.gather(params, indices, axis=axis)
 
 
 def _jax_fill(dims, value):
@@ -417,6 +593,127 @@ def _jax_transpose(a, perm=None):
   return jnp.transpose(a, axes=perm)
 
 
+def _jax_get_seed_data(seed: Any) -> Optional[np.ndarray]:
+  """Extracts the underlying numerical data from a JAX PRNGKey."""
+  if seed is None:
+    return None
+
+  return np.array(jax.random.key_data(seed))
+
+
+def _tf_get_seed_data(seed: Any) -> Optional[np.ndarray]:
+  """Converts a TensorFlow-style seed into a NumPy array."""
+  if seed is None:
+    return None
+  return np.array(seed)
+
+
+def _jax_convert_to_tensor(data, dtype=None):
+  """Converts data to a JAX array, handling strings as NumPy arrays."""
+  # JAX does not natively support string tensors in the same way TF does.
+  # If a string dtype is requested, or if the data is inherently strings,
+  # we fall back to a standard NumPy array.
+  is_string_target = False
+  if dtype is not None:
+    try:
+      if np.dtype(dtype).kind in ("S", "U"):
+        is_string_target = True
+    except TypeError:
+      # This can happen if dtype is not a valid dtype specifier,
+      # let jax.asarray handle it.
+      pass
+
+  is_string_data = isinstance(data, (list, np.ndarray)) and np.array(
+      data
+  ).dtype.kind in ("S", "U")
+
+  if is_string_target or (dtype is None and is_string_data):
+    return np.array(data, dtype=dtype)
+
+  return jax_ops.asarray(data, dtype=dtype)
+
+
+def _tf_nanmean(a, axis=None, keepdims=False):
+  import tensorflow.experimental.numpy as tnp
+
+  return tf_backend.convert_to_tensor(
+      tnp.nanmean(a, axis=axis, keepdims=keepdims)
+  )
+
+
+def _tf_nansum(a, axis=None, keepdims=False):
+  import tensorflow.experimental.numpy as tnp
+
+  return tf_backend.convert_to_tensor(
+      tnp.nansum(a, axis=axis, keepdims=keepdims)
+  )
+
+
+def _tf_nanvar(a, axis=None, keepdims=False):
+  """Calculates variance ignoring NaNs, strictly returning a Tensor."""
+  import tensorflow as tf
+  import tensorflow.experimental.numpy as tnp
+  # We implement two-pass variance to correctly handle NaNs and ensure
+  # all operations remain within the TF graph (maintaining differentiability).
+  a_tensor = tf.convert_to_tensor(a)
+  mean = tnp.nanmean(a_tensor, axis=axis, keepdims=True)
+  sq_diff = tf.math.squared_difference(a_tensor, mean)
+  var = tnp.nanmean(sq_diff, axis=axis, keepdims=keepdims)
+  return tf.convert_to_tensor(var)
+
+
+def _jax_one_hot(
+    indices, depth, on_value=None, off_value=None, axis=None, dtype=None
+):
+  """JAX implementation for one_hot."""
+  import jax.numpy as jnp
+
+  resolved_dtype = _resolve_dtype(dtype, on_value, off_value, 1, 0)
+  jax_axis = -1 if axis is None else axis
+
+  one_hot_result = jax.nn.one_hot(
+      indices, num_classes=depth, dtype=jnp.dtype(resolved_dtype), axis=jax_axis
+  )
+
+  on_val = 1 if on_value is None else on_value
+  off_val = 0 if off_value is None else off_value
+
+  if on_val == 1 and off_val == 0:
+    return one_hot_result
+
+  on_tensor = jnp.array(on_val, dtype=jnp.dtype(resolved_dtype))
+  off_tensor = jnp.array(off_val, dtype=jnp.dtype(resolved_dtype))
+
+  return jnp.where(one_hot_result == 1, on_tensor, off_tensor)
+
+
+def _jax_roll(a, shift, axis=None):
+  """JAX implementation for roll."""
+  import jax.numpy as jnp
+
+  return jnp.roll(a, shift, axis=axis)
+
+
+def _tf_roll(a, shift: Sequence[int], axis=None):
+  """TensorFlow implementation for roll that handles axis=None."""
+  import tensorflow as tf
+
+  if axis is None:
+    original_shape = tf.shape(a)
+    flat_tensor = tf.reshape(a, [-1])
+    rolled_flat = tf.roll(flat_tensor, shift=shift, axis=0)
+    return tf.reshape(rolled_flat, original_shape)
+  return tf.roll(a, shift, axis=axis)
+
+
+def _jax_enable_op_determinism():
+  """No-op for JAX. Determinism is handled via stateless PRNGKeys."""
+  warnings.warn(
+      "op determinism is a TensorFlow-specific concept and has no effect when"
+      " using the JAX backend."
+  )
+
+
 # --- Backend Initialization ---
 _BACKEND = config.get_backend()
 
@@ -429,15 +726,67 @@ if _BACKEND == config.Backend.JAX:
   import jax
   import jax.numpy as jax_ops
   import tensorflow_probability.substrates.jax as tfp_jax
+  from jax import tree_util
 
   class ExtensionType:
-    """A JAX-compatible stand-in for tf.experimental.ExtensionType."""
+    """A JAX-compatible stand-in for tf.experimental.ExtensionType.
 
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-      raise NotImplementedError(
-          "ExtensionType is not yet implemented for the JAX backend."
+    This class registers itself as a JAX Pytree node, allowing it to be passed
+    through JIT-compiled functions.
+    """
+
+    def __init_subclass__(cls, **kwargs):
+      super().__init_subclass__(**kwargs)
+      tree_util.register_pytree_node(
+          cls,
+          cls._tree_flatten,
+          cls._tree_unflatten,
       )
 
+    def _tree_flatten(self):
+      """Flattens the object for JAX tracing.
+
+      Fields containing JAX arrays (or convertibles) are treated as children.
+      Fields containing strings or NumPy string arrays must be treated as
+      auxiliary data because JAX cannot trace non-numeric/non-boolean data.
+
+      Returns:
+        A tuple of (children, aux_data), where children are the
+        tracer-compatible parts of the object, and aux_data contains auxiliary
+        information needed for unflattening.
+      """
+      d = vars(self)
+      all_keys = sorted(d.keys())
+      children = []
+      aux = {}
+      children_keys = []
+
+      for k in all_keys:
+        v = d[k]
+        # Identify string data to prevent JAX tracing errors.
+        # 'S' is zero-terminated bytes (fixed-width), 'U' is unicode string.
+        is_numpy_string = isinstance(v, np.ndarray) and v.dtype.kind in (
+            "S",
+            "U",
+        )
+        is_plain_string = isinstance(v, str)
+
+        if is_numpy_string or is_plain_string:
+          aux[k] = v
+        else:
+          children.append(v)
+          children_keys.append(k)
+      return children, (aux, children_keys)
+
+    @classmethod
+    def _tree_unflatten(cls, aux_and_keys, children):
+      aux, children_keys = aux_and_keys
+      obj = cls.__new__(cls)
+      vars(obj).update(aux)
+      for k, v in zip(children_keys, children):
+        setattr(obj, k, v)
+      return obj
+
   class _JaxErrors:
     # pylint: disable=invalid-name
     ResourceExhaustedError = MemoryError
@@ -492,6 +841,26 @@ if _BACKEND == config.Backend.JAX:
 
   random = _JaxRandom()
 
+  @functools.partial(
+      jax.jit,
+      static_argnames=[
+          "joint_dist",
+          "n_chains",
+          "n_draws",
+          "num_adaptation_steps",
+          "dual_averaging_kwargs",
+          "max_tree_depth",
+          "unrolled_leapfrog_steps",
+          "parallel_iterations",
+      ],
+  )
+  def _jax_xla_windowed_adaptive_nuts(**kwargs):
+    """JAX-specific JIT wrapper for the NUTS sampler."""
+    kwargs["seed"] = random.prng_key(kwargs["seed"])
+    return experimental.mcmc.windowed_adaptive_nuts(**kwargs)
+
+  xla_windowed_adaptive_nuts = _jax_xla_windowed_adaptive_nuts
+
   _ops = jax_ops
   errors = _JaxErrors()
   Tensor = jax.Array
@@ -499,7 +868,7 @@ if _BACKEND == config.Backend.JAX:
   bijectors = tfp_jax.bijectors
   experimental = tfp_jax.experimental
   mcmc = tfp_jax.mcmc
-  _convert_to_tensor = _ops.asarray
+  _convert_to_tensor = _jax_convert_to_tensor
 
   # Standardized Public API
   absolute = _ops.abs
@@ -507,18 +876,6 @@ if _BACKEND == config.Backend.JAX:
   arange = _jax_arange
   argmax = _jax_argmax
   boolean_mask = _jax_boolean_mask
-  concatenate = _ops.concatenate
-  stack = _ops.stack
-  split = _jax_split
-  zeros = _ops.zeros
-  zeros_like = _ops.zeros_like
-  ones = _ops.ones
-  ones_like = _ops.ones_like
-  repeat = _ops.repeat
-  reshape = _ops.reshape
-  tile = _ops.tile
-  where = _ops.where
-  broadcast_to = _ops.broadcast_to
   broadcast_dynamic_shape = _jax_broadcast_dynamic_shape
   broadcast_to = _ops.broadcast_to
   cast = _jax_cast
@@ -527,6 +884,7 @@ if _BACKEND == config.Backend.JAX:
   divide = _ops.divide
   divide_no_nan = _jax_divide_no_nan
   einsum = _ops.einsum
+  enable_op_determinism = _jax_enable_op_determinism
   equal = _ops.equal
   exp = _ops.exp
   expand_dims = _ops.expand_dims
@@ -534,12 +892,17 @@ if _BACKEND == config.Backend.JAX:
   function = _jax_function_wrapper
   gather = _jax_gather
   get_indices_where = _jax_get_indices_where
+  get_seed_data = _jax_get_seed_data
   is_nan = _ops.isnan
   log = _ops.log
   make_ndarray = _jax_make_ndarray
   make_tensor_proto = _jax_make_tensor_proto
+  nanmean = jax_ops.nanmean
   nanmedian = _jax_nanmedian
+  nansum = jax_ops.nansum
+  nanvar = jax_ops.nanvar
   numpy_function = _jax_numpy_function
+  one_hot = _jax_one_hot
   ones = _ops.ones
   ones_like = _ops.ones_like
   rank = _ops.ndim
@@ -551,6 +914,8 @@ if _BACKEND == config.Backend.JAX:
   reduce_sum = _ops.sum
   repeat = _ops.repeat
   reshape = _ops.reshape
+  roll = _jax_roll
+  split = _jax_split
   stack = _ops.stack
   tile = _jax_tile
   transpose = _jax_transpose
@@ -564,13 +929,14 @@ if _BACKEND == config.Backend.JAX:
   newaxis = _ops.newaxis
   TensorShape = _jax_tensor_shape
   int32 = _ops.int32
+  string = np.bytes_
+
+  stabilize_rf_roi_grid = _jax_stabilize_rf_roi_grid
 
   def set_random_seed(seed: int) -> None:  # pylint: disable=unused-argument
-    raise NotImplementedError(
-        "JAX does not support a global, stateful random seed. `set_random_seed`"
-        " is not implemented. Instead, you must pass an explicit `seed`"
-        " integer directly to the sampling methods (e.g., `sample_prior`),"
-        " which will be used to create a JAX PRNGKey internally."
+    warnings.warn(
+        "backend.set_random_seed is a no-op in JAX. Randomness is managed "
+        "statelessly via PRNGKeys passed to sampling methods."
     )
 
 elif _BACKEND == config.Backend.TENSORFLOW:
@@ -638,29 +1004,25 @@ elif _BACKEND == config.Backend.TENSORFLOW:
 
   random = _TfRandom()
 
+  @tf_backend.function(autograph=False, jit_compile=True)
+  def _tf_xla_windowed_adaptive_nuts(**kwargs):
+    """TensorFlow-specific XLA wrapper for the NUTS sampler."""
+    return experimental.mcmc.windowed_adaptive_nuts(**kwargs)
+
+  xla_windowed_adaptive_nuts = _tf_xla_windowed_adaptive_nuts
+
   tfd = tfp.distributions
   bijectors = tfp.bijectors
   experimental = tfp.experimental
   mcmc = tfp.mcmc
-
   _convert_to_tensor = _ops.convert_to_tensor
+
+  # Standardized Public API
   absolute = _ops.math.abs
   allclose = _ops.experimental.numpy.allclose
   arange = _tf_arange
   argmax = _tf_argmax
   boolean_mask = _tf_boolean_mask
-  concatenate = _ops.concat
-  stack = _ops.stack
-  split = _ops.split
-  zeros = _ops.zeros
-  zeros_like = _ops.zeros_like
-  ones = _ops.ones
-  ones_like = _ops.ones_like
-  repeat = _ops.repeat
-  reshape = _ops.reshape
-  tile = _ops.tile
-  where = _ops.where
-  broadcast_to = _ops.broadcast_to
   broadcast_dynamic_shape = _ops.broadcast_dynamic_shape
   broadcast_to = _ops.broadcast_to
   cast = _ops.cast
@@ -669,6 +1031,7 @@ elif _BACKEND == config.Backend.TENSORFLOW:
   divide = _ops.divide
   divide_no_nan = _ops.math.divide_no_nan
   einsum = _ops.einsum
+  enable_op_determinism = _ops.config.experimental.enable_op_determinism
   equal = _ops.equal
   exp = _ops.math.exp
   expand_dims = _ops.expand_dims
@@ -676,12 +1039,17 @@ elif _BACKEND == config.Backend.TENSORFLOW:
   function = _tf_function_wrapper
   gather = _tf_gather
   get_indices_where = _tf_get_indices_where
+  get_seed_data = _tf_get_seed_data
   is_nan = _ops.math.is_nan
   log = _ops.math.log
   make_ndarray = _ops.make_ndarray
   make_tensor_proto = _ops.make_tensor_proto
+  nanmean = _tf_nanmean
   nanmedian = _tf_nanmedian
+  nansum = _tf_nansum
+  nanvar = _tf_nanvar
   numpy_function = _ops.numpy_function
+  one_hot = _ops.one_hot
   ones = _ops.ones
   ones_like = _ops.ones_like
   rank = _ops.rank
@@ -693,7 +1061,9 @@ elif _BACKEND == config.Backend.TENSORFLOW:
   reduce_sum = _ops.reduce_sum
   repeat = _ops.repeat
   reshape = _ops.reshape
+  roll = _tf_roll
   set_random_seed = tf_backend.keras.utils.set_random_seed
+  split = _ops.split
   stack = _ops.stack
   tile = _ops.tile
   transpose = _ops.transpose
@@ -702,11 +1072,14 @@ elif _BACKEND == config.Backend.TENSORFLOW:
   zeros = _ops.zeros
   zeros_like = _ops.zeros_like
 
+  stabilize_rf_roi_grid = _tf_stabilize_rf_roi_grid
+
   float32 = _ops.float32
   bool_ = _ops.bool
   newaxis = _ops.newaxis
   TensorShape = _ops.TensorShape
   int32 = _ops.int32
+  string = _ops.string
 
 else:
   raise ValueError(f"Unsupported backend: {_BACKEND}")
@@ -735,7 +1108,7 @@ def _extract_int_seed(s: Any) -> Optional[int]:
   # in unpredictable ways. The goal is to safely attempt extraction and fail
   # gracefully.
   except Exception:  # pylint: disable=broad-except
-    pass
+    return None
   return None
 
 
@@ -819,6 +1192,14 @@ class _JaxRNGHandler(_BaseRNGHandler):
     if seed is None:
       return
 
+    if (
+        isinstance(seed, jax.Array)  # pylint: disable=undefined-variable
+        and seed.shape == (2,)
+        and seed.dtype == jax_ops.uint32  # pylint: disable=undefined-variable
+    ):
+      self._key = seed
+      return
+
     if self._int_seed is None:
       raise ValueError(
           "JAX backend requires a seed that is an integer or a scalar array,"
@@ -830,7 +1211,8 @@ class _JaxRNGHandler(_BaseRNGHandler):
   def get_kernel_seed(self) -> Any:
     if self._key is None:
       return None
-    return random.sanitize_seed(self._key)
+    _, subkey = random.split(self._key)
+    return int(jax.random.randint(subkey, (), 0, 2**31 - 1))  # pylint: disable=undefined-variable
 
   def get_next_seed(self) -> Any:
     if self._key is None:
@@ -853,35 +1235,22 @@ class _JaxRNGHandler(_BaseRNGHandler):
     return _JaxRNGHandler(np.asarray(new_seed_tensor).item())
 
 
-# TODO: Replace with _TFRNGHandler
-class _TFLegacyRNGHandler(_BaseRNGHandler):
-  """TensorFlow implementation.
+class _TFRNGHandler(_BaseRNGHandler):
+  """A stateless-style RNG handler for TensorFlow.
 
-  TODO: This class should be removed and replaced with a correct,
-  stateful `tf.random.Generator`-based implementation.
+  This handler canonicalizes any seed input into a single stateless seed Tensor.
   """
 
-  def __init__(self, seed: SeedType, *, _sanitized_seed: Optional[Any] = None):
-    """Initializes the TensorFlow legacy RNG handler.
+  def __init__(self, seed: SeedType):
+    """Initializes the TensorFlow RNG handler.
 
     Args:
       seed: The initial seed. Can be an integer, a sequence of two integers, a
-        corresponding Tensor, or None.
-      _sanitized_seed: For internal use only. If provided, this pre-computed
-        seed tensor is used directly, and the standard initialization logic for
-        the public `seed` argument is bypassed.
-
-    Raises:
-      ValueError: If `seed` is a sequence with a length other than 2.
+        corresponding Tensor, or None. It will be sanitized and stored
+        internally as a single stateless seed Tensor.
     """
     super().__init__(seed)
-    self._tf_sanitized_seed: Optional[Any] = None
-
-    if _sanitized_seed is not None:
-      # Internal path: A pre-sanitized seed was provided by a trusted source
-      # so we adopt it directly.
-      self._tf_sanitized_seed = _sanitized_seed
-      return
+    self._seed_state: Optional["_tf.Tensor"] = None
 
     if seed is None:
       return
@@ -899,63 +1268,33 @@ class _TFLegacyRNGHandler(_BaseRNGHandler):
     else:
       seed_to_sanitize = seed
 
-    self._tf_sanitized_seed = random.sanitize_seed(seed_to_sanitize)
+    self._seed_state = random.sanitize_seed(seed_to_sanitize)
 
   def get_kernel_seed(self) -> Any:
-    return self._tf_sanitized_seed
+    """Provides the current stateless seed of the handler without advancing it."""
+    return self._seed_state
 
   def get_next_seed(self) -> Any:
-    """Returns the original integer seed to preserve prior sampling behavior.
-
-    Returns:
-      The original integer seed provided during initialization.
-
-    Raises:
-      RuntimeError: If the handler was not initialized with a scalar integer
-        seed, which is required for the legacy prior sampling path.
-    """
-    if self._seed_input is None:
+    """Returns a new unique seed and advances the handler's internal state."""
+    if self._seed_state is None:
       return None
+    new_seed, next_state = random.stateless_split(self._seed_state)
+    self._seed_state = next_state
+    return new_seed
 
-    if self._int_seed is None:
-      raise RuntimeError(
-          "RNGHandler was not initialized with a scalar integer seed, cannot"
-          " provide seed for TensorFlow prior sampling."
-      )
-    return self._int_seed
-
-  def advance_handler(self) -> "_TFLegacyRNGHandler":
-    """Creates a new handler by incrementing the sanitized seed by 1.
-
-    Returns:
-      A new `_TFLegacyRNGHandler` instance with an incremented seed state.
-
-    Raises:
-      RuntimeError: If the handler's sanitized seed was not initialized.
-    """
-    if self._seed_input is None:
-      return _TFLegacyRNGHandler(None)
-
-    if self._tf_sanitized_seed is None:
-      # Should be caught during init, but included for defensive programming.
-      raise RuntimeError("RNGHandler sanitized seed not initialized.")
-
-    new_sanitized_seed = self._tf_sanitized_seed + 1
-
-    # Create a new handler instance, passing the original seed input (to
-    # preserve state like `_int_seed`) and injecting the new sanitized seed
-    # via the private constructor argument.
-    return _TFLegacyRNGHandler(
-        self._seed_input, _sanitized_seed=new_sanitized_seed
-    )
+  def advance_handler(self) -> "_TFRNGHandler":
+    """Creates a new handler from a new seed and advances the current handler."""
+    if self._seed_state is None:
+      return _TFRNGHandler(None)
+    seed_for_new_handler, next_state = random.stateless_split(self._seed_state)
+    self._seed_state = next_state
+    return _TFRNGHandler(seed_for_new_handler)
 
 
 if _BACKEND == config.Backend.JAX:
   RNGHandler = _JaxRNGHandler
 elif _BACKEND == config.Backend.TENSORFLOW:
-  RNGHandler = (
-      _TFLegacyRNGHandler  # TODO: Replace with _TFRNGHandler
-  )
+  RNGHandler = _TFRNGHandler
 else:
   raise ImportError(f"RNGHandler not implemented for backend: {_BACKEND}")