keras/src/backend/jax/numpy.py/expand_dims
def expand_dims(x, axis):
    x = convert_to_tensor(x)
    if isinstance(x, jax_sparse.BCOO):
        (
            _,
            result_shape,
            broadcast_dimensions,
        ) = sparse.axis_shape_dims_for_broadcast_in_dim(
            axis, x.shape, insert_dims=True
        )
        return jax_sparse.bcoo_broadcast_in_dim(
            x, shape=result_shape, broadcast_dimensions=broadcast_dimensions
        )
    return jnp.expand_dims(x, axis)