SciREX

Scientific research and engineering excellence

Period: 2024 – present
Affiliation: AiREX Lab, IISc · ZenteiQ · open source (Apache 2.0)

Open source
JAX
Scientific ML

SciREX (Scientific Research and Engineering eXcellence) is an open-source scientific AI and machine-learning framework for researchers and engineers, developed jointly by ZenteiQ Aitech Innovations and the AiREX Lab at the Indian Institute of Science, Bangalore. It aims at the gap between a method as published and a method you can actually run, without giving up mathematical rigour or computational efficiency.

Design goals

Research-first. Built around scientific computing and research workflows rather than retrofitted from a production ML stack.
Mathematically honest. Emphasis on correctness and theoretical grounding over convenience.
Fast where it matters. Efficient implementations with GPU acceleration.
Reproducible. Experiment tracking and result reproduction are part of the framework, not an afterthought.
Publication-ready output. Scientific visualisation built in.

Modules

Diffusion — noise schedules and diffusion samplers.
Training — utilities for optimisation loops and model training.
Transformers — modular Transformer and U-Net architectures.
Experimental — research-focused and in-flight implementations.

Using it

pip install scirex

import jax.numpy as jnp
from flax import nnx as nn
from scirex.training import Trainer

model = nn.Linear(10, 2, rngs=nn.Rngs(0))

def loss_fn(model, batch):
    x, y = batch
    logits = model(x)
    return jnp.mean((logits - y) ** 2), logits

trainer = Trainer(
    model=model,
    optimizer=nn.Optimizer(model, optax.adam(1e-3)),
    loss_fn=loss_fn,
)
trainer.train(batch_iterator, num_epochs=5)

Apache 2.0 licensed. Source on GitHub, documentation at scirex.org.