Remote GPUs for JAX,without the heavy rewrite.
Sahasra is built for teams that want cloud GPU access without turning adoption into an infrastructure project. Your JAX code stays familiar, the remote step gets simpler, and the advanced path is still there when you need it.
What improved
- `SahasraRuntime.connect(...)` now provides a cleaner entry point for runtime setup.
- `sahasra.jit(...)` gives new users a much smaller jump from normal JAX code.
- The public demo now also includes a higher-level remote training example.
Best for
Faster onboarding
Ideal when you want to show remote GPU value first, without making every user learn the low-level session model on day one.
Still possible
Explicit control
Teams that want direct lifecycle control can still use `SahasraClient` and `RemoteJaxSession`.
No local GPU
required on the client machine
2 API paths
high-level first, advanced when needed
Public demo
available at demo.sahasra.dev
March 18, 2026
latest API simplification update
Why teams look at Sahasra
A remote execution layer that is easier to adopt and easier to explain.
Keep your JAX code familiar
Sahasra lets you keep your function local, then route the heavy execution step to a remote GPU with a much cleaner wrapper than before.
Start with the simpler API
New users can begin with `SahasraRuntime.connect(...)` and `sahasra.jit(...)` instead of managing every session detail up front.
Drop lower only when you need to
The original `SahasraClient` and `RemoteJaxSession` APIs still exist for teams that want explicit control over remote execution and tensor handling.
Built for real remote workflows
The current public flow has already been validated with external clients talking to a remote worker, so this is moving toward real use, not just a concept page.
Quickstart
Start with the high-level path.
The recommended first experience is now `SahasraRuntime.connect(...)` plus `sahasra.jit(...)`. That keeps the function body familiar, reduces visible setup, and makes the remote story much easier for new users to understand.
Typical onboarding flow
- Install `sahasra` and `jax` in the client environment.
- Point the client at the public API URL and bearer token.
- Wrap a normal JAX function with `sahasra.jit(...)` and run it remotely.
import jax.numpy as jnp
import sahasra
runtime = sahasra.SahasraRuntime.connect(gpu_class="g5", region="ap-south-1")
@sahasra.jit(runtime=runtime)
def step(x):
return jnp.tanh(x @ x)
x = jnp.ones((128, 128), dtype=jnp.float32)
estimate = step.estimate(x)
y = step(x)
print(estimate.estimated_execute_sec)
print(y.shape, y.dtype)
runtime.close()Advanced mode
Drop lower when your workflow needs it.
Sahasra still supports the explicit client-and-session path for teams that want lower-level control. The point of the new API is not to remove power, but to stop making power the first thing every user has to see.
import jax
import jax.numpy as jnp
from sahasra import RemoteJaxSession, SahasraClient
client = SahasraClient(
base_url="https://demo.sahasra.dev",
api_key="REPLACE_WITH_CLIENT_TOKEN",
)
session = RemoteJaxSession.create(client, gpu_class="g5", region="ap-south-1")
@jax.jit
def step(x):
return jnp.tanh(x @ x)
execution = session.execute_jitted(
step,
sample_inputs=(jnp.ones((128, 128), dtype=jnp.float32),),
)
print(execution.runtime_mode)
session.close()Where Sahasra fits best right now
Teams testing GPU demand before heavy infra work
If your team writes JAX locally but does not want to stand up more infrastructure on day one, Sahasra gives you a clearer way to validate the need first.
Developers who want less ceremony
The high-level API is for users who want the remote step to feel close to normal JAX usage instead of becoming a session-management project.
Power users who still want control
If your workflow needs lower-level client setup, explicit estimates, or remote tensor handles, the original advanced path is still available.
Current proof point
The public MNIST demo has already been updated to use the new high-level API path, and the current verified remote training flow reaches about 93.5% test accuracy in that demo setup.
Contact
Want remote GPU access that is easier to show to users and easier to adopt internally?
Sahasra is being shaped into a cleaner front door for remote JAX execution: less setup for new users, an advanced path for power users, and a product story that makes sense without exposing every runtime detail.