Given the challenges of maintaining installations of various frameworks in a single environment, users who would want to test startai with multiple frameworks at once can use our Docker images for a seamless experience. You can pull the images from:

docker pull khulnasoft/startai:latest      # CPU
docker pull khulnasoft/startai:latest-gpu  # GPU

You can also install Startai from source if you want to take advantage of the latest changes, but we can't ensure everything will work as expected 😅

git clone https://github.com/khulnasoft/startai.git
cd startai
pip install --user -e .

If you want to set up testing and various frameworks it's probably best to check out the Setting Up page, where OS-specific and IDE-specific instructions and video tutorials to do so are available!

I'm using PyTorch 

You can use Startai to get PyTorch code from:

Any model

From TensorFlow

import startai
import torch
import tensorflow as tf

# Get a pretrained keras model
eff_encoder = tf.keras.applications.efficientnet_v2.EfficientNetV2B0(
    include_top=False, weights="imagenet", input_shape=(224, 224, 3)
)

# Transpile it into a torch.nn.Module with the corresponding parameters
noise = tf.random.normal(shape=(1, 224, 224, 3))
torch_eff_encoder = startai.transpile(eff_encoder, source="tensorflow", to="torch", args=(noise,))

# Build a classifier using the transpiled encoder
class Classifier(torch.nn.Module):
    def __init__(self, num_classes=20):
        super().__init__()
        self.encoder = torch_eff_encoder
        self.fc = torch.nn.Linear(1280, num_classes)

    def forward(self, x):
        x = self.encoder(x)
        return self.fc(x)

# Initialize a trainable, customizable, torch.nn.Module
classifier = Classifier()
ret = classifier(torch.rand((1, 244, 244, 3)))

From JAX

import startai
import jax
import torch

# Get a pretrained haiku model
# https://github.com/khulnasoft/demos/blob/15c235f/scripts/deepmind_perceiver_io.py
from deepmind_perceiver_io import key, perceiver_backbone

# Transpile it into a torch.nn.Module with the corresponding parameters
dummy_input = jax.random.uniform(key, shape=(1, 3, 224, 224))
params = perceiver_backbone.init(rng=key, images=dummy_input)
startai.set_backend("jax")
backbone = startai.transpile(
    perceiver_backbone, source="jax", to="torch", params_v=params, kwargs={"images": dummy_input}
)

# Build a classifier using the transpiled backbone
class PerceiverIOClassifier(torch.nn.Module):
    def __init__(self, num_classes=20):
        super().__init__()
        self.backbone = backbone
        self.max_pool = torch.nn.MaxPool2d((512, 1))
        self.flatten = torch.nn.Flatten()
        self.fc = torch.nn.Linear(1024, num_classes)

    def forward(self, x):
        x = self.backbone(images=x)
        x = self.flatten(self.max_pool(x))
        return self.fc(x)

# Initialize a trainable, customizable, torch.nn.Module
classifier = PerceiverIOClassifier()
ret = classifier(torch.rand((1, 3, 224, 224)))

Any library

From Tensorflow

import startai
import torch
import os
os.environ["SM_FRAMEWORK"] = "tf.keras"
import segmentation_models as sm

# transpile sm from tensorflow to torch
torch_sm = startai.transpile(sm, source="tensorflow", to="torch")

# get some image-like arrays
output = torch.rand((1, 3, 512, 512))
target = torch.rand((1, 3, 512, 512))

# and use the transpiled version of any function from the library!
out = torch_sm.metrics.iou_score(output, target)

From JAX

import startai
import rax
import torch

# transpile rax from jax to torch
torch_rax = startai.transpile(rax, source="jax", to="torch")

# get some arrays
scores = torch.tensor([2.2, 1.3, 5.4])
labels = torch.tensor([1.0, 0.0, 0.0])

# and use the transpiled version of any function from the library!
out = torch_rax.poly1_softmax_loss(scores, labels)

From NumPy

import startai
import torch
import madmom

# transpile madmon from numpy to torch
torch_madmom = startai.transpile(madmom, source="numpy", to="torch")

# get some arrays
freqs = torch.arange(20) * 10

# and use the transpiled version of any function from the library!
out = torch_madmom.audio.filters.hz2midi(freqs)

Any function

From Tensorflow

import startai
import tensorflow as tf
import torch

def loss(predictions, targets):
    return tf.sqrt(tf.reduce_mean(tf.square(predictions - targets)))

# transpile any function from tf to torch
torch_loss = startai.transpile(loss, source="tensorflow", to="torch")

# get some arrays
p = torch.tensor([3.0, 2.0, 1.0])
t = torch.tensor([0.0, 0.0, 0.0])

# and use the transpiled version!
out = torch_loss(p, t)

From JAX

import startai
import jax.numpy as jnp
import torch

def loss(predictions, targets):
    return jnp.sqrt(jnp.mean((predictions - targets) ** 2))

# transpile any function from jax to torch
torch_loss = startai.transpile(loss, source="jax", to="torch")

# get some arrays
p = torch.tensor([3.0, 2.0, 1.0])
t = torch.tensor([0.0, 0.0, 0.0])

# and use the transpiled version!
out = torch_loss(p, t)

From NumPy

import startai
import numpy as np
import torch

def loss(predictions, targets):
    return np.sqrt(np.mean((predictions - targets) ** 2))

# transpile any function from numpy to torch
torch_loss = startai.transpile(loss, source="numpy", to="torch")

# get some arrays
p = torch.tensor([3.0, 2.0, 1.0])
t = torch.tensor([0.0, 0.0, 0.0])

# and use the transpiled version!
out = torch_loss(p, t)

I'm using TensorFlow 

You can use Startai to get TensorFlow code from:

Any model

From PyTorch

import startai
import torch
import timm
import tensorflow as tf

# Get a pretrained pytorch model
mlp_encoder = timm.create_model("mixer_b16_224", pretrained=True, num_classes=0)

# Transpile it into a keras.Model with the corresponding parameters
noise = torch.randn(1, 3, 224, 224)
mlp_encoder = startai.transpile(mlp_encoder, to="tensorflow", args=(noise,))

# Build a classifier using the transpiled encoder
class Classifier(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.encoder = mlp_encoder
        self.output_dense = tf.keras.layers.Dense(units=1000, activation="softmax")

    def call(self, x):
        x = self.encoder(x)
        return self.output_dense(x)

# Transform the classifier and use it as a standard keras.Model
x = tf.random.normal(shape=(1, 3, 224, 224))
model = Classifier()
ret = model(x)

From JAX

import startai
import jax
import tensorflow as tf

# Get a pretrained haiku model
# https://khulnasoft.com/demos/scripts/deepmind_perceiver_io.py
from deepmind_perceiver_io import key, perceiver_backbone

# Transpile it into a tf.keras.Model with the corresponding parameters
dummy_input = jax.random.uniform(key, shape=(1, 3, 224, 224))
params = perceiver_backbone.init(rng=key, images=dummy_input)
backbone = startai.transpile(
    perceiver_backbone, to="tensorflow", params_v=params, args=(dummy_input,)
)

# Build a classifier using the transpiled backbone
class PerceiverIOClassifier(tf.keras.Model):
    def __init__(self, num_classes=20):
        super().__init__()
        self.backbone = backbone
        self.max_pool = tf.keras.layers.MaxPooling1D(pool_size=512)
        self.flatten = tf.keras.layers.Flatten()
        self.fc = tf.keras.layers.Dense(num_classes)

    def call(self, x):
        x = self.backbone(x)
        x = self.flatten(self.max_pool(x))
        return self.fc(x)

# Initialize a trainable, customizable, tf.keras.Model
x = tf.random.normal(shape=(1, 3, 224, 224))
classifier = PerceiverIOClassifier()
ret = classifier(x)

Any library

From PyTorch

import startai
import kornia
import requests
import numpy as np
import tensorflow as tf
from PIL import Image

# transpile kornia from torch to tensorflow
tf_kornia = startai.transpile(kornia, source="torch", to="tensorflow")

# get an image
url = "http://images.cocodataset.org/train2017/000000000034.jpg"
raw_img = Image.open(requests.get(url, stream=True).raw)

# convert it to the format expected by kornia
img = np.array(raw_img)
img = tf.transpose(tf.constant(img), (2, 0, 1))
img = tf.expand_dims(img, 0) / 255

# and use the transpiled version of any function from the library!
out = tf_kornia.enhance.sharpness(img, 5)

From JAX

import startai
import rax
import tensorflow as tf

# transpile rax from jax to tensorflow
tf_rax = startai.transpile(rax, source="jax", to="tensorflow")

# get some arrays
scores = tf.constant([2.2, 1.3, 5.4])
labels = tf.constant([1.0, 0.0, 0.0])

# and use the transpiled version of any function from the library!
out = tf_rax.poly1_softmax_loss(scores, labels)

From NumPy

import startai
import madmom
import tensorflow as tf

# transpile madmom from numpy to tensorflow
tf_madmom = startai.transpile(madmom, source="numpy", to="tensorflow")

# get some arrays
freqs = tf.range(20) * 10

# and use the transpiled version of any function from the library!
out = tf_madmom.audio.filters.hz2midi(freqs)

Any function

From PyTorch

import startai
import torch
import tensorflow as tf

def loss(predictions, targets):
    return torch.sqrt(torch.mean((predictions - targets) ** 2))

# transpile any function from torch to tensorflow
tf_loss = startai.transpile(loss, source="torch", to="tensorflow")

# get some arrays
p = tf.constant([3.0, 2.0, 1.0])
t = tf.constant([0.0, 0.0, 0.0])

# and use the transpiled version!
out = tf_loss(p, t)

From JAX

import startai
import jax.numpy as jnp
import tensorflow as tf

def loss(predictions, targets):
    return jnp.sqrt(jnp.mean((predictions - targets) ** 2))

# transpile any function from jax to tensorflow
tf_loss = startai.transpile(loss, source="jax", to="tensorflow")

# get some arrays
p = tf.constant([3.0, 2.0, 1.0])
t = tf.constant([0.0, 0.0, 0.0])

# and use the transpiled version!
out = tf_loss(p, t)

From NumPy

import startai
import numpy as np
import tensorflow as tf

def loss(predictions, targets):
    return np.sqrt(np.mean((predictions - targets) ** 2))

# transpile any function from numpy to tensorflow
tf_loss = startai.transpile(loss, source="numpy", to="tensorflow")

# get some arrays
p = tf.constant([3.0, 2.0, 1.0])
t = tf.constant([0.0, 0.0, 0.0])

# and use the transpiled version!
out = tf_loss(p, t)

I'm using Jax 

You can use Startai to get JAX code from:

Any model

From PyTorch

import startai
import timm
import torch
import jax
import haiku as hk

# Get a pretrained pytorch model
mlp_encoder = timm.create_model("mixer_b16_224", pretrained=True, num_classes=0)

# Transpile it into a hk.Module with the corresponding parameters
noise = torch.randn(1, 3, 224, 224)
mlp_encoder = startai.transpile(mlp_encoder, source="torch", to="haiku", args=(noise,))

# Build a classifier using the transpiled encoder
class Classifier(hk.Module):
    def __init__(self, num_classes=1000):
        super().__init__()
        self.encoder = mlp_encoder()
        self.fc = hk.Linear(output_size=num_classes, with_bias=True)

    def __call__(self, x):
        x = self.encoder(x)
        x = self.fc(x)
        return x

def _forward_classifier(x):
    module = Classifier()
    return module(x)

# Transform the classifier and use it as a standard hk.Module
rng_key = jax.random.PRNGKey(42)
x = jax.random.uniform(key=rng_key, shape=(1, 3, 224, 224), dtype=jax.numpy.float32)
forward_classifier = hk.transform(_forward_classifier)
params = forward_classifier.init(rng=rng_key, x=x)

ret = forward_classifier.apply(params, None, x)

From TensorFlow

import startai
import jax
import haiku as hk
import tensorflow as tf
jax.config.update("jax_enable_x64", True)

# Get a pretrained keras model
eff_encoder = tf.keras.applications.efficientnet_v2.EfficientNetV2B0(
    include_top=False, weights="imagenet", input_shape=(224, 224, 3)
)

# Transpile it into a hk.Module with the corresponding parameters
noise = tf.random.normal(shape=(1, 224, 224, 3))
hk_eff_encoder = startai.transpile(eff_encoder, source="tensorflow", to="haiku", args=(noise,))

# Build a classifier using the transpiled encoder
class Classifier(hk.Module):
    def __init__(self, num_classes=1000):
        super().__init__()
        self.encoder = hk_eff_encoder()
        self.fc = hk.Linear(output_size=num_classes, with_bias=True)

    def __call__(self, x):
        x = self.encoder(x)
        x = self.fc(x)
        return x

def _forward_classifier(x):
    module = Classifier()
    return module(x)

# Transform the classifier and use it as a standard hk.Module
rng_key = jax.random.PRNGKey(42)
dummy_x = jax.random.uniform(key=rng_key, shape=(1, 224, 224, 3))
forward_classifier = hk.transform(_forward_classifier)
params = forward_classifier.init(rng=rng_key, x=dummy_x)

ret = forward_classifier.apply(params, None, dummy_x)

Any library

From PyTorch

import startai
import kornia
import requests
import jax.numpy as jnp
from PIL import Image
jax.config.update("jax_enable_x64", True)

# transpile kornia from torch to jax
jax_kornia = startai.transpile(kornia, source="torch", to="jax")

# get an image
url = "http://images.cocodataset.org/train2017/000000000034.jpg"
raw_img = Image.open(requests.get(url, stream=True).raw)

# convert it to the format expected by kornia
img = jnp.transpose(jnp.array(raw_img), (2, 0, 1))
img = jnp.expand_dims(img, 0) / 255

# and use the transpiled version of any function from the library!
out = jax_kornia.enhance.sharpness(img, 5)

From TensorFlow

import startai
import jax
import os
os.environ["SM_FRAMEWORK"] = "tf.keras"
import segmentation_models as sm

# transpile sm from tensorflow to jax
jax_sm = startai.transpile(sm, source="tensorflow", to="jax")

# get some image-like arrays
key = jax.random.PRNGKey(23)
key1, key2 = jax.random.split(key)
output = jax.random.uniform(key1, (1, 3, 512, 512))
target = jax.random.uniform(key2, (1, 3, 512, 512))

# and use the transpiled version of any function from the library!
out = jax_sm.metrics.iou_score(output, target)

From NumPy

import startai
import madmom
import jax.numpy as jnp

# transpile madmon from numpy to jax
jax_madmom = startai.transpile(madmom, source="numpy", to="jax")

# get some arrays
freqs = jnp.arange(20) * 10

# and use the transpiled version of any function from the library!
out = jax_madmom.audio.filters.hz2midi(freqs)

Any function

From PyTorch

import startai
import torch
import jax.numpy as jnp

def loss(predictions, targets):
    return torch.sqrt(torch.mean((predictions - targets) ** 2))

# transpile any function from torch to jax
jax_loss = startai.transpile(loss, source="torch", to="jax")

# get some arrays
p = jnp.array([3.0, 2.0, 1.0])
t = jnp.array([0.0, 0.0, 0.0])

# and use the transpiled version!
out = jax_loss(p, t)

From TensorFlow

import startai
import tensorflow as tf
import jax.numpy as jnp

def loss(predictions, targets):
    return tf.sqrt(tf.reduce_mean(tf.square(predictions - targets)))

# transpile any function from tf to jax
jax_loss = startai.transpile(loss, source="tensorflow", to="jax")

# get some arrays
p = jnp.array([3.0, 2.0, 1.0])
t = jnp.array([0.0, 0.0, 0.0])

# and use the transpiled version!
out = jax_loss(p, t)

From NumPy

import startai
import numpy as np
import jax
import jax.numpy as jnp
jax.config.update('jax_enable_x64', True)

def loss(predictions, targets):
    return np.sqrt(np.mean((predictions - targets) ** 2))

# transpile any function from numpy to jax
jax_loss = startai.transpile(loss, source="numpy", to="jax")

# get some arrays
p = jnp.array([3.0, 2.0, 1.0])
t = jnp.array([0.0, 0.0, 0.0])

# and use the transpiled version!
out = jax_loss(p, t)

I'm using NumPy 

You can use Startai to get NumPy code from:

Any library

From PyTorch

import startai
import kornia
import requests
import numpy as np
from PIL import Image

# transpile kornia from torch to np
np_kornia = startai.transpile(kornia, source="torch", to="numpy")

# get an image
url = "http://images.cocodataset.org/train2017/000000000034.jpg"
raw_img = Image.open(requests.get(url, stream=True).raw)

# convert it to the format expected by kornia
img = np.transpose(np.array(raw_img), (2, 0, 1))
img = np.expand_dims(img, 0) / 255

# and use the transpiled version of any function from the library!
out = np_kornia.enhance.sharpness(img, 5)

From TensorFlow

import startai
import numpy as np
import os
os.environ["SM_FRAMEWORK"] = "tf.keras"
import segmentation_models as sm

# transpile sm from tensorflow to numpy
np_sm = startai.transpile(sm, source="tensorflow", to="numpy")

# get some image-like arrays
output = np.random.rand(1, 3, 512, 512).astype(dtype=np.float32)
target = np.random.rand(1, 3, 512, 512).astype(dtype=np.float32)

# and use the transpiled version of any function from the library!
out = np_sm.metrics.iou_score(output, target)

From Jax

import startai
import rax
import numpy as np

# transpile rax from jax to numpy
np_rax = startai.transpile(rax, source="jax", to="numpy")

# get some arrays
scores = np.array([2.2, 1.3, 5.4])
labels = np.array([1.0, 0.0, 0.0])

# and use the transpiled version of any function from the library!
out = np_rax.poly1_softmax_loss(scores, labels)

Any function

From PyTorch

import startai
import torch
import numpy as np

def loss(predictions, targets):
    return torch.sqrt(torch.mean((predictions - targets) ** 2))

# transpile any function from torch to numpy
np_loss = startai.transpile(loss, source="torch", to="numpy")

# get some arrays
p = np.array([3.0, 2.0, 1.0])
t = np.array([0.0, 0.0, 0.0])

# and use the transpiled version!
out = np_loss(p, t)

From TensorFlow

import startai
import tensorflow as tf
import numpy as np

def loss(predictions, targets):
    return tf.sqrt(tf.reduce_mean(tf.square(predictions - targets)))

# transpile any function from tf to numpy
np_loss = startai.transpile(loss, source="tensorflow", to="numpy")

# get some arrays
p = np.array([3.0, 2.0, 1.0])
t = np.array([0.0, 0.0, 0.0])

# and use the transpiled version!
out = np_loss(p, t)

From JAX

import startai
import jax.numpy as jnp
import numpy as np

def loss(predictions, targets):
    return jnp.sqrt(jnp.mean((predictions - targets) ** 2))

# transpile any function from jax to numpy
np_loss = startai.transpile(loss, source="jax", to="numpy")

# get some arrays
p = np.array([3.0, 2.0, 1.0])
t = np.array([0.0, 0.0, 0.0])

# and use the transpiled version!
out = np_loss(p, t)

I'm using Startai 

import startai

# A simple image classification model
class StartaiNet(startai.Module):
    def __init__(
        self,
        h_w=(32, 32),
        input_channels=3,
        output_channels=512,
        num_classes=2,
        data_format="NCHW",
        device="cpu",
    ):
        self.h_w = h_w
        self.input_channels = input_channels
        self.output_channels = output_channels
        self.num_classes = num_classes
        self.data_format = data_format
        super().__init__(device=device)

    def _build(self, *args, **kwargs):
        self.extractor = startai.Sequential(
            startai.Conv2D(self.input_channels, 6, [5, 5], 1, "SAME", data_format=self.data_format),
            startai.GELU(),
            startai.Conv2D(6, 16, [5, 5], 1, "SAME", data_format=self.data_format),
            startai.GELU(),
            startai.Conv2D(16, self.output_channels, [5, 5], 1, "SAME", data_format=self.data_format),
            startai.GELU(),
        )

        self.classifier = startai.Sequential(
            # Since the padding is "SAME", this would be image_height x image_width x output_channels
            startai.Linear(self.h_w[0] * self.h_w[1] * self.output_channels, 512),
            startai.GELU(),
            startai.Linear(512, self.num_classes),
        )

    def _forward(self, x):
        x = self.extractor(x)
        # flatten all dims except batch dim
        x = startai.flatten(x, start_dim=1, end_dim=-1)
        logits = self.classifier(x)
        probs = startai.softmax(logits)
        return logits, probs

After building your model in Startai, you can set your favourite framework as the backend to use its operations under the hood!

startai.set_backend("torch")
model = StartaiNet()
x = torch.randn(1, 3, 32, 32)
logits, probs = model(x)

startai.set_backend("tensorflow")
model = StartaiNet()
x = tf.random.uniform(shape=(1, 3, 32, 32))
logits, probs = model(x)

startai.set_backend("jax")
model = StartaiNet()
x = jax.random.uniform(key, shape=(1, 3, 32, 32))
logits, probs = model(x)

startai.set_backend("numpy")
model = StartaiNet()
x = np.random.uniform(size=(1, 3, 32, 32))
logits, probs = model(x)

Last but not least, we can also build the training pipeline in pure startai ⬇️

# helper function for loading the dataset in batches
def generate_batches(images, classes, dataset_size, batch_size=32):
    if batch_size > dataset_size:
        raise startai.utils.exceptions.StartaiError("Use a smaller batch size")
    for idx in range(0, dataset_size, batch_size):
        yield images[idx : min(idx + batch_size, dataset_size)], classes[
            idx : min(idx + batch_size, dataset_size)
        ]


# helper function to get the number of current predictions
def num_correct(preds, labels):
    return (preds.argmax() == labels).sum().to_numpy().item()


# define a loss function
def loss_fn(params):
    v, model, x, y = params
    _, probs = model(x, v=v)
    return startai.cross_entropy(y, probs), probs

# train the model on gpu if it's available
device = "gpu:0" if startai.gpu_is_available() else "cpu"

# training hyperparams
optimizer = startai.Adam(1e-4)
batch_size = 4
num_epochs = 20
num_classes = 10

model = StartaiNet(
    h_w=(28, 28),
    input_channels=1,
    output_channels=120,
    num_classes=num_classes,
    device=device,
)

images = startai.random_uniform(shape=(16, 1, 28, 28))
classes = startai.randint(0, num_classes - 1, shape=(16,))


# training loop
def train(images, classes, epochs, model, device, num_classes=10, batch_size=32):
    # training metrics
    epoch_loss = 0.0
    metrics = []
    dataset_size = len(images)

    for epoch in range(epochs):
        train_correct = 0
        train_loop = tqdm(
            generate_batches(images, classes, len(images), batch_size=batch_size),
            total=dataset_size // batch_size,
            position=0,
            leave=True,
        )

        for xbatch, ybatch in train_loop:
            xbatch, ybatch = xbatch.to_device(device), ybatch.to_device(device)

            # Since the cross entropy function expects the target classes to be in one-hot encoded format
            ybatch_encoded = startai.one_hot(ybatch, num_classes)

            # update model params
            loss_probs, grads = startai.execute_with_gradients(
                loss_fn,
                (model.v, model, xbatch, ybatch_encoded),
            )

            model.v = optimizer.step(model.v, grads["0"])

            batch_loss = startai.to_numpy(loss_probs[0]).mean().item()  # batch mean loss
            epoch_loss += batch_loss * xbatch.shape[0]
            train_correct += num_correct(loss_probs[1], ybatch)

            train_loop.set_description(f"Epoch [{epoch + 1:2d}/{epochs}]")
            train_loop.set_postfix(
                running_loss=batch_loss,
                accuracy_percentage=(train_correct / dataset_size) * 100,
            )

        epoch_loss = epoch_loss / dataset_size
        training_accuracy = train_correct / dataset_size

        metrics.append([epoch, epoch_loss, training_accuracy])

        train_loop.write(
            f"\nAverage training loss: {epoch_loss:.6f}, Train Correct: {train_correct}",
            end="\n",
        )


# assuming the dataset(images and classes) are already prepared in a folder
train(
    images,
    classes,
    num_epochs,
    model,
    device,
    num_classes=num_classes,
    batch_size=batch_size,
)

When should I use Startai as a transpiler?

If you want to use building blocks published in other frameworks (neural networks, layers, array computing libraries, training pipelines...), you want to integrate code developed in various frameworks, or maybe straight up move code from one framework to another, the transpiler is definitely the tool 🔧 for the job! As the output of transpilation is native code in the target framework, you can use the converted code just as if it was code originally developed in that framework, applying framework-specific optimizations or tools, instantly exposing your project to all of the unique perks of a different framework.

Startai's transpiler allows you to use code from any other framework (or from any other version of the same framework!) in your own code, by just adding one line of code. Under the hood, Startai traces a computational graph and leverages the frontends and backends to link one framework to another.

This way, Startai makes all ML-related projects available for you, independently of the framework you want to use to research, develop, or deploy systems. Feel free to head over to the docs for the full API reference, but the functions you'd most likely want to use are:

# Traces an efficient fully-functional graph from a function, removing all wrapping and redundant code
startai.trace_graph()

# Converts framework-specific code to a different framework
startai.transpile()

# Converts framework-specific code to Startai
startai.unify()

These functions can be used eagerly or lazily. If you pass the necessary arguments for function tracing, the graph tracing/transpilation step will happen instantly (eagerly). Otherwise, the graph tracing/transpilation will happen only when the returned function is first invoked.

import startai
import jax
startai.set_backend("jax")

# Simple JAX function to transpile
def test_fn(x):
    return jax.numpy.sum(x)

x1 = startai.array([1., 2.])

# Arguments are available -> transpilation happens eagerly
eager_graph = startai.transpile(test_fn, source="jax", to="torch", args=(x1,))

# eager_graph is now torch code and runs efficiently
ret = eager_graph(x1)

# Arguments are not available -> transpilation happens lazily
lazy_graph = startai.transpile(test_fn, source="jax", to="torch")

# The transpiled graph is initialized, transpilation will happen here
ret = lazy_graph(x1)

# lazy_graph is now torch code and runs efficiently
ret = lazy_graph(x1)

If you want to learn more, you can find more information in the Startai as a transpiler section of the docs!

When should I use Startai as a framework?

As Startai supports multiple backends, writing code in Startai breaks you free from framework limitations. If you want to publish highly flexible code for everyone to use, independently of the framework they are using, or you plan to develop ML-related tools and want them to be interoperable with not only the already existing frameworks, but also with future frameworks, then Startai is for you!

The Startai framework is built on top of various essential components, mainly the Backend Handler, which manages what framework is being used behind the scenes and the Backend Functional APIs, which provide framework-specific implementations of the Startai functions. Likewise, classes such as startai.Container or startai.Array are also available, facilitating the use of structured data and array-like objects (learn more about them here!).

All of the functionalities in Startai are exposed through the Startai functional API and the Startai stateful API. All functions in the Functional API are Framework Agnostic Functions, which means that we can use them like this:

import startai
import jax.numpy as jnp
import tensorflow as tf
import numpy as np
import torch

def mse_loss(y, target):
    return startai.mean((y - target)**2)

jax_mse   = mse_loss(jnp.ones((5,)), jnp.ones((5,)))
tf_mse    = mse_loss(tf.ones((5,)), tf.ones((5,)))
np_mse    = mse_loss(np.ones((5,)), np.ones((5,)))
torch_mse = mse_loss(torch.ones((5,)), torch.ones((5,)))

In the example above we show how Startai's functions are compatible with tensors from different frameworks. This is the same for ALL Startai functions. They can accept tensors from any framework and return the correct result.

The Startai Stateful API, on the other hand, allows you to define trainable modules and layers, which you can use alone or as a part of any other framework code!

import startai


class Regressor(startai.Module):
    def __init__(self, input_dim, output_dim):
        self.input_dim = input_dim
        self.output_dim = output_dim
        super().__init__()

    def _build(self, *args, **kwargs):
        self.linear0 = startai.Linear(self.input_dim, 128)
        self.linear1 = startai.Linear(128, self.output_dim)

    def _forward(self, x):
        x = self.linear0(x)
        x = startai.functional.relu(x)
        x = self.linear1(x)
        return x

If we put it all together, we'll have something like this. This example uses PyTorch as the backend, but this can easily be changed to your favorite frameworks, such as TensorFlow, or JAX.

import startai


class Regressor(startai.Module):
    def __init__(self, input_dim, output_dim):
        self.input_dim = input_dim
        self.output_dim = output_dim
        super().__init__()

    def _build(self, *args, **kwargs):
        self.linear0 = startai.Linear(self.input_dim, 128)
        self.linear1 = startai.Linear(128, self.output_dim)

    def _forward(self, x):
        x = self.linear0(x)
        x = startai.functional.relu(x)
        x = self.linear1(x)
        return x

startai.set_backend('torch')  # set backend to PyTorch (or any other backend!)

model = Regressor(input_dim=1, output_dim=1)
optimizer = startai.Adam(0.3)

n_training_examples = 2000
noise = startai.random.random_normal(shape=(n_training_examples, 1), mean=0, std=0.1)
x = startai.linspace(-6, 3, n_training_examples).reshape((n_training_examples, 1))
y = 0.2 * x ** 2 + 0.5 * x + 0.1 + noise


def loss_fn(v, x, target):
    pred = model(x, v=v)
    return startai.mean((pred - target) ** 2)

for epoch in range(40):
    # forward pass
    pred = model(x)

    # compute loss and gradients
    loss, grads = startai.execute_with_gradients(lambda params: loss_fn(*params), (model.v, x, y))

    # update parameters
    model.v = optimizer.step(model.v, grads)

    # print current loss
    print(f'Epoch: {epoch + 1:2d} --- Loss: {startai.to_numpy(loss).item():.5f}')

print('Finished training!')

The model's output can be visualized as follows:

As always, you can find more information about Startai as a framework in the docs!