phiml 1.0.3

Unified API for machine learning

Φ_ML

🌐 Homepage   •   📖 Documentation   •   🔗 API   •   ▶ Videos   •   Examples

Φ_ML provides a unified math and neural network API for Jax, PyTorch, TensorFlow and NumPy.

See the installation Instructions on how to compile the optional custom CUDA operations.

from jax import numpy as jnp
import torch
import tensorflow as tf
import numpy as np

from phiml import math

math.sin(1.)
math.sin(jnp.asarray([1.]))
math.sin(torch.tensor([1.]))
math.sin(tf.constant([1.]))
math.sin(np.asarray([1.]))

Compatibility

• Writing code that works with PyTorch, Jax, and TensorFlow makes it easier to share code with other people and collaborate.
• Your published research code will reach a broader audience.
• When you run into a bug / roadblock with one library, you can simply switch to another.
• Φ_ML can efficiently convert tensors between ML libraries on-the-fly, so you can even mix the different ecosystems.

Fewer mistakes

• No more data type troubles: Φ_ML automatically converts data types where needed and lets you specify the FP precision globally or by context!
• No more reshaping troubles: Φ_ML performs reshaping under-the-hood.
• Is neighbor_idx.at[jnp.reshape(idx, (-1,))].set(jnp.reshape(cell_idx, (-1,) + cell_idx.shape[-2:])) correct?: Φ_ML provides a custom Tensor class that lets you write easy-to-read, more concise, more explicit, less error-prone code.

Unique features

• n-dimensional operations: With Φ_ML, you can write code that automatically works in 1D, 2D and 3D, choosing the corresponding operations based on the input dimensions.
• Preconditioned linear solves: Φ_ML can build sparse matrices from your Python functions and run linear solvers with preconditioners.
• Flexible neural network architectures: Φ_ML provides various configurable neural network architectures, from MLPs to U-Nets.
• Non-uniform tensors: Φ_ML allows you to stack tensors of different sizes and keeps track of the resulting shapes.