google-meridian 1.1.5__py3-none-any.whl → 1.2.0__py3-none-any.whl

meridian/backend/__init__.py

@@ -0,0 +1,514 @@
+# Copyright 2025 The Meridian Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Backend Abstraction Layer for Meridian."""
+
+import os
+from typing import Any, Optional, TYPE_CHECKING, Tuple, Union
+
+from meridian.backend import config
+import numpy as np
+from typing_extensions import Literal
+
+# The conditional imports in this module are a deliberate design choice for the
+# backend abstraction layer. The TFP-on-JAX substrate provides a nearly
+# identical API to the standard TFP library, making an alias-based approach more
+# pragmatic than a full Abstract Base Class implementation, which would require
+# extensive boilerplate.
+# pylint: disable=g-import-not-at-top,g-bad-import-order
+
+if TYPE_CHECKING:
+  import dataclasses
+  import jax as _jax
+  import tensorflow as _tf
+
+  TensorShapeInstance = Union[_tf.TensorShape, Tuple[int, ...]]
+
+
+def standardize_dtype(dtype: Any) -> str:
+  """Converts a backend-specific dtype to a standard string representation.
+
+  Args:
+    dtype: A backend-specific dtype object (e.g., tf.DType, np.dtype).
+
+  Returns:
+    A canonical string representation of the dtype (e.g., 'float32').
+  """
+
+  # Handle None explicitly, as np.dtype(None) defaults to float64.
+
+  if dtype is None:
+    return str(None)
+
+  if hasattr(dtype, "as_numpy_dtype"):
+    dtype = dtype.as_numpy_dtype
+
+  try:
+    return np.dtype(dtype).name
+  except TypeError:
+    return str(dtype)
+
+
+def result_type(*types: Any) -> str:
+  """Infers the result dtype from a list of input types, backend-agnostically.
+
+  This acts as the single source of truth for type promotion rules. The
+  promotion logic is designed to be consistent across all backends.
+
+  Rule: If any input is a float, the result is float32. Otherwise, the result
+  is int64 to match NumPy/JAX's default behavior for precision.
+
+  Args:
+    *types: A variable number of type objects (e.g., `<class 'int'>`,
+      np.dtype('float32')).
+
+  Returns:
+    A string representing the promoted dtype.
+  """
+  standardized_types = []
+  for t in types:
+    if t is None:
+      continue
+    try:
+      # Standardize the input type before checking promotion rules.
+      standardized_types.append(standardize_dtype(t))
+    except Exception:  # pylint: disable=broad-except
+      # Fallback if standardization fails for an unexpected type.
+      standardized_types.append(str(t))
+
+  if any("float" in t for t in standardized_types):
+    return "float32"
+  return "int64"
+
+
+def _resolve_dtype(dtype: Optional[Any], *args: Any) -> str:
+  """Resolves the final dtype for an operation.
+
+  If a dtype is explicitly provided, it's returned. Otherwise, it infers the
+  dtype from the input arguments using the backend-agnostic `result_type`
+  promotion rules.
+
+  Args:
+    dtype: The user-provided dtype, which may be None.
+    *args: The input arguments to the operation, used for dtype inference.
+
+  Returns:
+    A string representing the resolved dtype.
+  """
+  if dtype is not None:
+    return standardize_dtype(dtype)
+
+  input_types = [
+      getattr(arg, "dtype", type(arg)) for arg in args if arg is not None
+  ]
+  return result_type(*input_types)
+
+
+# --- Private Backend-Specific Implementations ---
+
+
+def _jax_arange(
+    start: Any,
+    stop: Optional[Any] = None,
+    step: Any = 1,
+    dtype: Optional[Any] = None,
+) -> "_jax.Array":
+  """JAX implementation for arange."""
+
+  # Import locally to make the function self-contained.
+
+  import jax.numpy as jnp
+
+  resolved_dtype = _resolve_dtype(dtype, start, stop, step)
+  return jnp.arange(start, stop, step=step, dtype=resolved_dtype)
+
+
+def _tf_arange(
+    start: Any,
+    stop: Optional[Any] = None,
+    step: Any = 1,
+    dtype: Optional[Any] = None,
+) -> "_tf.Tensor":
+  """TensorFlow implementation for arange."""
+  import tensorflow as tf
+
+  resolved_dtype = _resolve_dtype(dtype, start, stop, step)
+  try:
+    return tf.range(start, limit=stop, delta=step, dtype=resolved_dtype)
+  except tf.errors.NotFoundError:
+    result = tf.range(start, limit=stop, delta=step, dtype=tf.float32)
+    return tf.cast(result, resolved_dtype)
+
+
+def _jax_cast(x: Any, dtype: Any) -> "_jax.Array":
+  """JAX implementation for cast."""
+  return x.astype(dtype)
+
+
+def _jax_divide_no_nan(x, y):
+  """JAX implementation for divide_no_nan."""
+  import jax.numpy as jnp
+
+  return jnp.where(y != 0, jnp.divide(x, y), 0.0)
+
+
+def _jax_numpy_function(*args, **kwargs):  # pylint: disable=unused-argument
+  raise NotImplementedError(
+      "backend.numpy_function is not implemented for the JAX backend."
+  )
+
+
+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_ndarray(*args, **kwargs):  # pylint: disable=unused-argument
+  raise NotImplementedError(
+      "backend.make_ndarray is not implemented for the JAX backend."
+  )
+
+
+def _jax_get_indices_where(condition):
+  """JAX implementation for get_indices_where."""
+  import jax.numpy as jnp
+
+  return jnp.stack(jnp.where(condition), axis=-1)
+
+
+def _tf_get_indices_where(condition):
+  """TensorFlow implementation for get_indices_where."""
+  import tensorflow as tf
+
+  return tf.where(condition)
+
+
+def _jax_unique_with_counts(x):
+  """JAX implementation for unique_with_counts."""
+  import jax.numpy as jnp
+
+  y, counts = jnp.unique(x, return_counts=True)
+  # The TF version returns a tuple of (y, idx, count). The idx is not used in
+  # the calling code, so we can return None for it to maintain tuple structure.
+  return y, None, counts
+
+
+def _tf_unique_with_counts(x):
+  """TensorFlow implementation for unique_with_counts."""
+  import tensorflow as tf
+
+  return tf.unique_with_counts(x)
+
+
+def _jax_boolean_mask(tensor, mask, axis=None):
+  """JAX implementation for boolean_mask that supports an axis argument."""
+  import jax.numpy as jnp
+
+  if axis is None:
+    axis = 0
+  tensor_swapped = jnp.moveaxis(tensor, axis, 0)
+  masked = tensor_swapped[mask]
+  return jnp.moveaxis(masked, 0, axis)
+
+
+def _tf_boolean_mask(tensor, mask, axis=None):
+  """TensorFlow implementation for boolean_mask."""
+  import tensorflow as tf
+
+  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 _tf_gather(params, indices):
+  """TensorFlow implementation for gather."""
+  import tensorflow as tf
+
+  return tf.gather(params, indices)
+
+
+def _jax_fill(dims, value):
+  """JAX implementation for fill."""
+  import jax.numpy as jnp
+
+  return jnp.full(dims, value)
+
+
+def _tf_fill(dims, value):
+  """TensorFlow implementation for fill."""
+  import tensorflow as tf
+
+  return tf.fill(dims, value)
+
+
+def _jax_argmax(tensor, axis=None):
+  """JAX implementation for argmax, aligned with TensorFlow's default.
+
+  This function finds the indices of the maximum values along a specified axis.
+  Crucially, it mimics the default behavior of TensorFlow's `tf.argmax`, where
+  if `axis` is `None`, the operation defaults to `axis=0`. This differs from
+  NumPy's and JAX's native `argmax` behavior, which would flatten the array
+  before finding the index.
+
+  Args:
+    tensor: The input JAX array.
+    axis: An integer specifying the axis along which to find the index of the
+      maximum value. If `None`, it defaults to `0` to match TensorFlow's
+      behavior.
+
+  Returns:
+    A JAX array containing the indices of the maximum values.
+  """
+  import jax.numpy as jnp
+
+  if axis is None:
+    axis = 0
+  return jnp.argmax(tensor, axis=axis)
+
+
+def _tf_argmax(tensor, axis=None):
+  """TensorFlow implementation for argmax."""
+  import tensorflow as tf
+
+  return tf.argmax(tensor, axis=axis)
+
+
+def _jax_broadcast_dynamic_shape(shape_x, shape_y):
+  """JAX implementation for broadcast_dynamic_shape."""
+  import jax.numpy as jnp
+
+  return jnp.broadcast_shapes(shape_x, shape_y)
+
+
+def _jax_tensor_shape(dims):
+  """JAX implementation for TensorShape."""
+  if isinstance(dims, int):
+    return (dims,)
+
+  return tuple(dims)
+
+
+# --- Backend Initialization ---
+_BACKEND = config.get_backend()
+
+# We expose standardized functions directly at the module level (backend.foo)
+# to provide a consistent, NumPy-like API across backends. The '_ops' object
+# is a private member for accessing the full, raw backend library if necessary,
+# but usage should prefer the top-level standardized functions.
+
+if _BACKEND == config.Backend.JAX:
+  import jax
+  import jax.numpy as jax_ops
+  import tensorflow_probability.substrates.jax as tfp_jax
+
+  class 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."
+      )
+
+  class _JaxErrors:
+    # pylint: disable=invalid-name
+    ResourceExhaustedError = MemoryError
+    InvalidArgumentError = ValueError
+    # pylint: enable=invalid-name
+
+  _ops = jax_ops
+  errors = _JaxErrors()
+  Tensor = jax.Array
+  tfd = tfp_jax.distributions
+  bijectors = tfp_jax.bijectors
+  experimental = tfp_jax.experimental
+  random = tfp_jax.random
+  mcmc = tfp_jax.mcmc
+  _convert_to_tensor = _ops.asarray
+
+  # Standardized Public API
+  absolute = _ops.abs
+  allclose = _ops.allclose
+  arange = _jax_arange
+  argmax = _jax_argmax
+  boolean_mask = _jax_boolean_mask
+  concatenate = _ops.concatenate
+  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
+  transpose = _ops.transpose
+  broadcast_to = _ops.broadcast_to
+  broadcast_dynamic_shape = _jax_broadcast_dynamic_shape
+  broadcast_to = _ops.broadcast_to
+  cast = _jax_cast
+  concatenate = _ops.concatenate
+  cumsum = _ops.cumsum
+  divide = _ops.divide
+  divide_no_nan = _jax_divide_no_nan
+  einsum = _ops.einsum
+  equal = _ops.equal
+  exp = _ops.exp
+  expand_dims = _ops.expand_dims
+  fill = _jax_fill
+  function = jax.jit
+  gather = _jax_gather
+  get_indices_where = _jax_get_indices_where
+  is_nan = _ops.isnan
+  log = _ops.log
+  make_ndarray = _jax_make_ndarray
+  make_tensor_proto = _jax_make_tensor_proto
+  numpy_function = _jax_numpy_function
+  ones = _ops.ones
+  ones_like = _ops.ones_like
+  rank = _ops.ndim
+  reduce_any = _ops.any
+  reduce_max = _ops.max
+  reduce_mean = _ops.mean
+  reduce_min = _ops.min
+  reduce_std = _ops.std
+  reduce_sum = _ops.sum
+  repeat = _ops.repeat
+  reshape = _ops.reshape
+  split = _ops.split
+  stack = _ops.stack
+  tile = _ops.tile
+  transpose = _ops.transpose
+  unique_with_counts = _jax_unique_with_counts
+  where = _ops.where
+  zeros = _ops.zeros
+  zeros_like = _ops.zeros_like
+
+  float32 = _ops.float32
+  bool_ = _ops.bool_
+  newaxis = _ops.newaxis
+  TensorShape = _jax_tensor_shape
+
+  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."
+    )
+
+elif _BACKEND == config.Backend.TENSORFLOW:
+  import tensorflow as tf_backend
+  import tensorflow_probability as tfp
+
+  _ops = tf_backend
+  errors = _ops.errors
+
+  Tensor = tf_backend.Tensor
+  ExtensionType = _ops.experimental.ExtensionType
+
+  tfd = tfp.distributions
+  bijectors = tfp.bijectors
+  experimental = tfp.experimental
+  random = tfp.random
+  mcmc = tfp.mcmc
+
+  _convert_to_tensor = _ops.convert_to_tensor
+  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
+  transpose = _ops.transpose
+  broadcast_to = _ops.broadcast_to
+  broadcast_dynamic_shape = _ops.broadcast_dynamic_shape
+  broadcast_to = _ops.broadcast_to
+  cast = _ops.cast
+  concatenate = _ops.concat
+  cumsum = _ops.cumsum
+  divide = _ops.divide
+  divide_no_nan = _ops.math.divide_no_nan
+  einsum = _ops.einsum
+  equal = _ops.equal
+  exp = _ops.math.exp
+  expand_dims = _ops.expand_dims
+  fill = _tf_fill
+  function = _ops.function
+  gather = _tf_gather
+  get_indices_where = _tf_get_indices_where
+  is_nan = _ops.math.is_nan
+  log = _ops.math.log
+  make_ndarray = _ops.make_ndarray
+  make_tensor_proto = _ops.make_tensor_proto
+  numpy_function = _ops.numpy_function
+  ones = _ops.ones
+  ones_like = _ops.ones_like
+  rank = _ops.rank
+  reduce_any = _ops.reduce_any
+  reduce_max = _ops.reduce_max
+  reduce_mean = _ops.reduce_mean
+  reduce_min = _ops.reduce_min
+  reduce_std = _ops.math.reduce_std
+  reduce_sum = _ops.reduce_sum
+  repeat = _ops.repeat
+  reshape = _ops.reshape
+  set_random_seed = tf_backend.keras.utils.set_random_seed
+  split = _ops.split
+  stack = _ops.stack
+  tile = _ops.tile
+  transpose = _ops.transpose
+  unique_with_counts = _tf_unique_with_counts
+  where = _ops.where
+  zeros = _ops.zeros
+  zeros_like = _ops.zeros_like
+
+  float32 = _ops.float32
+  bool_ = _ops.bool
+  newaxis = _ops.newaxis
+  TensorShape = _ops.TensorShape
+
+else:
+  raise ValueError(f"Unsupported backend: {_BACKEND}")
+# pylint: enable=g-import-not-at-top,g-bad-import-order
+
+
+def to_tensor(data: Any, dtype: Optional[Any] = None) -> Tensor:  # type: ignore
+  """Converts input data to the currently active backend tensor type.
+
+  Args:
+    data: The data to convert.
+    dtype: The desired data type of the resulting tensor. The accepted types
+      depend on the active backend (e.g., jax.numpy.dtype or tf.DType).
+
+  Returns:
+    A tensor representation of the data for the active backend.
+  """
+
+  return _convert_to_tensor(data, dtype=dtype)

meridian/backend/config.py

@@ -0,0 +1,59 @@
+# Copyright 2025 The Meridian Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Backend configuration for Meridian."""
+
+import enum
+import warnings
+
+
+class Backend(enum.Enum):
+  TENSORFLOW = "tensorflow"
+  JAX = "jax"
+
+
+_BACKEND = Backend.TENSORFLOW
+
+
+def set_backend(backend: Backend) -> None:
+  """Sets the backend for Meridian.
+
+  Note: The JAX backend is currently under development and should not be used.
+
+  Args:
+    backend: The backend to use, must be a member of the `Backend` enum.
+
+  Raises:
+    ValueError: If the provided backend is not a valid `Backend` enum member.
+  """
+  global _BACKEND
+  if not isinstance(backend, Backend):
+    raise ValueError("Backend must be a member of the Backend enum.")
+
+  if backend == Backend.JAX:
+    warnings.warn(
+        (
+            "The JAX backend is currently under development and is not yet"
+            " functional. It is intended for internal testing only and should"
+            " not be used. Please use the TensorFlow backend."
+        ),
+        UserWarning,
+    )
+
+  _BACKEND = backend
+
+
+def get_backend() -> Backend:
+  """Returns the current backend for Meridian."""
+  return _BACKEND

meridian/backend/test_utils.py

@@ -0,0 +1,95 @@
+# Copyright 2025 The Meridian Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Common testing utilities for Meridian, designed to be backend-agnostic."""
+
+from typing import Any
+import numpy as np
+
+# A type alias for backend-agnostic array-like objects.
+# We use `Any` here to avoid circular dependencies with the backend module
+# while still allowing the function to accept backend-specific tensor types.
+ArrayLike = Any
+
+
+def assert_allclose(
+    a: ArrayLike,
+    b: ArrayLike,
+    rtol: float = 1e-6,
+    atol: float = 1e-6,
+    err_msg: str = "",
+):
+  """Backend-agnostic assertion to check if two array-like objects are close.
+
+  This function converts both inputs to NumPy arrays before comparing them,
+  making it compatible with TensorFlow Tensors, JAX Arrays, and standard
+  Python lists or NumPy arrays.
+
+  Args:
+    a: The first array-like object to compare.
+    b: The second array-like object to compare.
+    rtol: The relative tolerance parameter.
+    atol: The absolute tolerance parameter.
+    err_msg: The error message to be printed in case of failure.
+
+  Raises:
+    AssertionError: If the two arrays are not equal within the given tolerance.
+  """
+  np.testing.assert_allclose(
+      np.array(a), np.array(b), rtol=rtol, atol=atol, err_msg=err_msg
+  )
+
+
+def assert_allequal(a: ArrayLike, b: ArrayLike, err_msg: str = ""):
+  """Backend-agnostic assertion to check if two array-like objects are equal.
+
+  This function converts both inputs to NumPy arrays before comparing them.
+
+  Args:
+    a: The first array-like object to compare.
+    b: The second array-like object to compare.
+    err_msg: The error message to be printed in case of failure.
+
+  Raises:
+    AssertionError: If the two arrays are not equal.
+  """
+  np.testing.assert_array_equal(np.array(a), np.array(b), err_msg=err_msg)
+
+
+def assert_all_finite(a: ArrayLike, err_msg: str = ""):
+  """Backend-agnostic assertion to check if all elements in an array are finite.
+
+  Args:
+    a: The array-like object to check.
+    err_msg: The error message to be printed in case of failure.
+
+  Raises:
+    AssertionError: If the array contains non-finite values.
+  """
+  if not np.all(np.isfinite(np.array(a))):
+    raise AssertionError(err_msg or "Array contains non-finite values.")
+
+
+def assert_all_non_negative(a: ArrayLike, err_msg: str = ""):
+  """Backend-agnostic assertion to check if all elements are non-negative.
+
+  Args:
+    a: The array-like object to check.
+    err_msg: The error message to be printed in case of failure.
+
+  Raises:
+    AssertionError: If the array contains negative values.
+  """
+  if not np.all(np.array(a) >= 0):
+    raise AssertionError(err_msg or "Array contains negative values.")