deepxde.backend.jax package

Submodules

deepxde.backend.jax.tensor module

jax backend implementation

deepxde.backend.jax.tensor.Variable(initial_value, dtype=None)

deepxde.backend.jax.tensor.abs(x)

deepxde.backend.jax.tensor.as_tensor(data, dtype=None)

deepxde.backend.jax.tensor.concat(values, axis)

deepxde.backend.jax.tensor.cos(x)

deepxde.backend.jax.tensor.data_type_dict()

deepxde.backend.jax.tensor.elu(x)

deepxde.backend.jax.tensor.from_numpy(np_array)

deepxde.backend.jax.tensor.is_tensor(obj)

deepxde.backend.jax.tensor.max(input_tensor, dim, keepdims=False)

deepxde.backend.jax.tensor.mean(input_tensor, dim, keepdims=False)

deepxde.backend.jax.tensor.min(input_tensor, dim, keepdims=False)

deepxde.backend.jax.tensor.minimum(x, y)

deepxde.backend.jax.tensor.ndim(input_array)

deepxde.backend.jax.tensor.prod(input_tensor, dim, keepdims=False)

deepxde.backend.jax.tensor.reduce_max(input_tensor)

deepxde.backend.jax.tensor.reduce_mean(input_tensor)

deepxde.backend.jax.tensor.reduce_min(input_tensor)

deepxde.backend.jax.tensor.reduce_prod(input_tensor)

deepxde.backend.jax.tensor.reduce_sum(input_tensor)

deepxde.backend.jax.tensor.relu(x)

deepxde.backend.jax.tensor.reshape(tensor, shape)

deepxde.backend.jax.tensor.selu(x)

deepxde.backend.jax.tensor.shape(input_array)

deepxde.backend.jax.tensor.sigmoid(x)

deepxde.backend.jax.tensor.silu(x)

deepxde.backend.jax.tensor.sin(x)

deepxde.backend.jax.tensor.square(x)

deepxde.backend.jax.tensor.stack(values, axis)

deepxde.backend.jax.tensor.sum(input_tensor, dim, keepdims=False)

deepxde.backend.jax.tensor.tanh(x)

deepxde.backend.jax.tensor.to_numpy(input_tensor)

deepxde.backend.jax.tensor.transpose(tensor, axes=None)

deepxde.backend.jax.tensor.zeros(shape, dtype)

deepxde.backend.jax.tensor.zeros_like(input_tensor)

Module contents