def _create_variables(self, values, trainable):
        """Create a structure of variables from a structure of JAX arrays.

        `values` is traversed via JAX's `tree_map`. When a leaf is a JAX array
        or a tensor-like object, a corresponding variable is created with it as
        the initial value. The resulting structure of variables is assigned to
        `self.params` or `self.state` depending on `trainable`. Then, a
        flattened version of the variables is returned for tracking.
        `self.params` or `self.state` are intentionally not tracked because
        structures like `TrackedList` interfere with `jax.tree_utils`.
        Note that leaf objects that are not JAX arrays and not tensor-like are
        left intact as they are assumed to be configuration used by the model.

        Args:
            values: the structure of values to traverse.
            trainable: whether to create trainable variables.

        Returns:
            flat list of variables initialized with `values` for tracking.
        """

        def create_variable(value):
            if backend.is_tensor(value) or isinstance(
                value, (np.ndarray, np.generic)
            ):
                dtype = value.dtype
                if is_float_dtype(dtype):
                    dtype = None  # Use the layer dtype policy
                return self.add_weight(
                    value.shape,
                    initializer=value,
                    dtype=dtype,
                    trainable=trainable,
                )
            elif isinstance(value, (bool, int, float)):
                dtype = standardize_dtype(type(value))
                if is_float_dtype(dtype):
                    dtype = None  # Use the layer dtype policy
                return self.add_weight(
                    (),
                    initializer=backend.convert_to_tensor(value),
                    dtype=dtype,
                    trainable=trainable,
                )
            else:
                return value

        # Use JAX's tree_map as it understands registered classes.
        variables = jax.tree_util.tree_map(create_variable, values)

        if trainable:
            self.params = variables
        else:
            self.state = variables

        flat_variables, _ = jax.tree_util.tree_flatten(variables)
        return flat_variables
Create a structure of variables from a structure of JAX arrays.

        `values` is traversed via JAX's `tree_map`. When a leaf is a JAX array
        or a tensor-like object, a corresponding variable is created with it as
        the initial value. The resulting structure of variables is assigned to
        `self.params` or `self.state` depending on `trainable`. Then, a
        flattened version of the variables is returned for tracking.
        `self.params` or `self.state` are intentionally not tracked because
        structures like `TrackedList` interfere with `jax.tree_utils`.
        Note that leaf objects that are not JAX arrays and not tensor-like are
        left intact as they are assumed to be configuration used by the model.

        Args:
            values: the structure of values to traverse.
            trainable: whether to create trainable variables.

        Returns:
            flat list of variables initialized with `values` for tracking.
        
keras/src/utils/jax_layer.py