API Reference

Setup

jask.set_memory_budget(max_memory, scratch_dir=None, allow_tmpfs=False)[source]

Set the process-wide configuration used by every jask op.

Call this once at the start of your program, not per-op-call. Every op (jask.dot, jask.add, etc.) that doesn’t receive an explicit Config uses this to decide how large a tile (“page”) of an array it reads into memory at a time, and where it writes its own internal files (op outputs, gradient buffers, spill files).

The effective memory budget is min(max_memory, 0.8 * currently available RAM) - a one-time safety clamp taken at the moment this function is called, so a budget that doesn’t match actual conditions on this machine doesn’t go completely unchecked. This is checked once here, not continuously during execution - block size stays fixed for the rest of the session once set. It does not protect against available memory dropping significantly after this call, only against a mismatch that already exists when you call it.

scratch_dir is checked against RAM-backed filesystems (tmpfs, ramfs - e.g. /tmp on many modern Linux systems and most containers) and rejected by default: writing “disk-backed” array data there would silently consume real memory instead of disk, defeating jask’s out-of-core guarantee entirely and invisibly.

Parameters:

max_memory (int or str) – The memory budget, either in bytes (int) or as a string with a unit suffix, e.g. "4GB", "512MB".
scratch_dir (str, optional) – Directory for jask’s own internal files. Defaults to .jask_scratch under the current working directory - still subject to the same RAM-backed-filesystem check, so a bad default is caught, not trusted blindly. Created if it doesn’t exist.
allow_tmpfs (bool, optional) – Set True to allow scratch_dir to be RAM-backed anyway (e.g. deliberately testing with small data). Default False.

Raises:

RuntimeError – If the resolved scratch_dir is on a RAM-backed filesystem and allow_tmpfs is False.

Examples

>>> import jask
>>> jask.set_memory_budget("4GB", scratch_dir="/data/jask_scratch")

The array type

class jask.DiskArray(filename, shape, dtype, _lo_tracer=None)[source]

A disk-backed array that behaves as a native JAX value.

DiskArray represents an array whose data lives entirely on disk and is streamed through JAX one tile at a time - never fully loaded into memory - while still composing with jax.grad, jax.jit, and optax exactly like an ordinary jax.Array would.

Parameters:

filename (str) – Path to the memmap file backing this array’s data.
shape (tuple of int) – The array’s shape.
dtype (numpy.dtype) – The array’s dtype.
_lo_tracer (object)

Notes

_lo_tracer is an internal field, set only transiently by DiskArrayType.raise_val or a primitive’s own expand, and is always a trivial marker (never real array data - see DiskArrayType for the zero-materialization mechanism). Real data always lives at filename. Not part of equality or repr.

Examples

>>> import numpy as np
>>> import jask
>>> jask.set_memory_budget("1GB")
>>> a = jask.DiskArray.from_numpy(np.ones((4, 4), dtype=np.float32))
>>> a.shape
(4, 4)

to_memmap()[source]

Return a read/write numpy.memmap view of this array’s file.

Returns:: A memmap of shape self.shape and dtype self.dtype. Reads and writes go straight to the backing file; no data is copied into memory until a specific region is accessed.
Return type:: numpy.memmap

classmethod from_numpy(arr)[source]

Write a numpy array to a temp file and wrap it as a DiskArray.

Intended for quick experiments with arrays small enough to build in memory first. For genuinely large data, construct DiskArray directly with an existing file’s path instead.

Parameters:: arr (numpy.ndarray) – The array to write to disk.
Returns:: A new disk-backed array holding a copy of arr’s data.
Return type:: DiskArray

Examples

>>> import numpy as np
>>> import jask
>>> a = jask.DiskArray.from_numpy(np.arange(4, dtype=np.float32))
>>> np.asarray(a.to_memmap())
array([0., 1., 2., 3.], dtype=float32)

update_(new_value)[source]

Overwrite this array’s file in place with new_value’s data.

Copies new_value’s data into this array’s own backing file, tile by tile (never a full in-memory copy), and returns a DiskArray with the same filename/identity as self.

This is the mechanism for loop-carried state (parameters, optimizer buffers) under jax.jit. Reassigning a = a + updates gives a a fresh filename every call, which makes jax.jit retrace on every single step (a new filename is a new type, so it’s a cache miss). a = a.update_(a + updates) keeps a’s filename identical across calls, so jax.jit compiles once and reuses that executable for the rest of the loop.

Parameters:: new_value (DiskArray) – The array whose data should replace this array’s contents. Must have the same shape and dtype as self.
Returns:: A DiskArray with the same filename as self, now holding new_value’s data.
Return type:: DiskArray

Examples

>>> import numpy as np
>>> import jask
>>> jask.set_memory_budget("1GB")
>>> a = jask.DiskArray.from_numpy(np.zeros((4, 4), dtype=np.float32))
>>> original_path = a.filename
>>> new_value = jask.DiskArray.from_numpy(np.ones((4, 4), dtype=np.float32))
>>> a = a.update_(new_value)
>>> a.filename == original_path
True
>>> np.asarray(a.to_memmap())[0, 0]
1.0

Operations

jask.add(*args, **kwargs)

Elementwise sum of two disk-backed arrays.

Computes a + b one tile at a time, never materializing either input or the output in full. Equivalent to a + b via DiskArray’s __add__.

Parameters:

a (DiskArray) – First operand.
b (DiskArray) – Second operand. Must have the same shape and dtype as a.

Returns:

A new disk-backed array of the same shape as a, holding a + b.

Return type:

DiskArray

Examples

>>> import numpy as np
>>> import jask
>>> jask.set_memory_budget("1GB")
>>> a = jask.DiskArray.from_numpy(np.ones((4, 4), dtype=np.float32))
>>> b = jask.DiskArray.from_numpy(np.ones((4, 4), dtype=np.float32))
>>> c = jask.add(a, b)
>>> np.asarray(c.to_memmap())[0, 0]
2.0

jask.sub(*args, **kwargs)

Elementwise difference of two disk-backed arrays.

Computes a - b one tile at a time, never materializing either input or the output in full. Equivalent to a - b via DiskArray’s __sub__.

Parameters:

a (DiskArray) – First operand (the minuend).
b (DiskArray) – Second operand (the subtrahend). Must have the same shape and dtype as a.

Returns:

A new disk-backed array of the same shape as a, holding a - b.

Return type:

DiskArray

Examples

>>> import numpy as np
>>> import jask
>>> jask.set_memory_budget("1GB")
>>> a = jask.DiskArray.from_numpy(np.full((4, 4), 5.0, dtype=np.float32))
>>> b = jask.DiskArray.from_numpy(np.ones((4, 4), dtype=np.float32))
>>> c = jask.sub(a, b)
>>> np.asarray(c.to_memmap())[0, 0]
4.0

jask.mul(a, b)[source]

Elementwise or scalar multiplication of disk-backed arrays.

Computes a * b one tile at a time, never materializing either input or the output in full. If either argument is not a DiskArray (a Python number or a jax scalar, concrete or traced - e.g. a learning rate passed as a jitted argument), the other is scaled by it elementwise. Order-independent: mul(a, 3.5) and mul(3.5, a) both work. Equivalent to a * b via DiskArray’s __mul__/__rmul__.

Parameters:

a (DiskArray or scalar) – First operand. At least one of a, b must be a DiskArray.
b (DiskArray or scalar) – Second operand.

Returns:

A new disk-backed array of the same shape as whichever operand is a DiskArray.

Return type:

DiskArray

Raises:

TypeError – If neither a nor b is a DiskArray.

Examples

>>> import numpy as np
>>> import jask
>>> jask.set_memory_budget("1GB")
>>> a = jask.DiskArray.from_numpy(np.full((4, 4), 2.0, dtype=np.float32))

Elementwise, both operands DiskArray:

>>> b = jask.DiskArray.from_numpy(np.full((4, 4), 3.0, dtype=np.float32))
>>> np.asarray(jask.mul(a, b).to_memmap())[0, 0]
6.0

Scalar, either order - the scalar can be a Python literal or a real (even jit-traced) jax value:

>>> np.asarray(jask.mul(a, 2.5).to_memmap())[0, 0]
5.0
>>> np.asarray(jask.mul(2.5, a).to_memmap())[0, 0]
5.0

jask.square(*args, **kwargs)

Elementwise square of a disk-backed array.

Computes a ** 2 one tile at a time, never materializing the input or output in full.

Parameters:: a (DiskArray) – The array to square.
Returns:: A new disk-backed array of the same shape as a, holding a ** 2.
Return type:: DiskArray

Examples

>>> import numpy as np
>>> import jask
>>> jask.set_memory_budget("1GB")
>>> a = jask.DiskArray.from_numpy(np.full((4, 4), 3.0, dtype=np.float32))
>>> b = jask.square(a)
>>> np.asarray(b.to_memmap())[0, 0]
9.0

jask.transpose(*args, **kwargs)

Permute the axes of a disk-backed array.

Computes the transpose one tile at a time, never materializing the input or output in full.

Parameters:

a (DiskArray) – The array to transpose.
axes (tuple of int, optional) – A permutation of range(a.ndim). If not given, reverses the order of all axes (numpy’s default transpose behavior).

Returns:

A new disk-backed array with axes permuted according to axes.

Return type:

DiskArray

Examples

>>> import numpy as np
>>> import jask
>>> jask.set_memory_budget("1GB")
>>> a = jask.DiskArray.from_numpy(np.arange(6, dtype=np.float32).reshape(2, 3))
>>> t = jask.transpose(a)
>>> np.asarray(t.to_memmap()).shape
(3, 2)

Permute a 3D array’s axes explicitly:

>>> b = jask.DiskArray.from_numpy(np.zeros((3, 4, 5), dtype=np.float32))
>>> jask.transpose(b, axes=(2, 0, 1)).shape
(5, 3, 4)

jask.sum(*args, **kwargs)

Sum all elements of a disk-backed array to a scalar.

Reduces the whole array one tile at a time, never materializing it in full. Unlike every other jask op, the result is small enough to return as a real jax.Array directly rather than a DiskArray.

Parameters:: a (DiskArray) – The array to reduce.
Returns:: A real (not disk-backed) scalar jax.Array holding the sum of every element of a.
Return type:: jax.Array

Examples

>>> import numpy as np
>>> import jask
>>> jask.set_memory_budget("1GB")
>>> a = jask.DiskArray.from_numpy(np.ones((4, 4), dtype=np.float32))
>>> float(jask.sum(a))
16.0

jask.dot(*args, **kwargs)

Matrix multiplication of two disk-backed arrays, with batching.

Computes a @ b one tile at a time along the contraction dimension, never materializing either input or the output in full. Matches jnp.matmul/@ semantics: for 2D inputs this is plain matrix multiplication; for higher-rank inputs, leading dimensions are treated as batch dimensions (broadcast, not contracted). Equivalent to a @ b via DiskArray’s __matmul__.

Note this differs from numpy.dot/jnp.dot’s own N-D contraction rule (which sums over a’s last axis and b’s second-to-last axis without batching) - jask.dot always matches @.

Parameters:

a (DiskArray) – Left operand, shape (*batch, M, K).
b (DiskArray) – Right operand, shape (*batch, K, N). Batch dimensions must match a’s exactly (no broadcasting between different batch shapes yet).

Returns:

A new disk-backed array of shape (*batch, M, N).

Return type:

DiskArray

Examples

>>> import numpy as np
>>> import jask
>>> jask.set_memory_budget("1GB")
>>> a = jask.DiskArray.from_numpy(np.ones((4, 6), dtype=np.float32))
>>> b = jask.DiskArray.from_numpy(np.ones((6, 3), dtype=np.float32))
>>> c = jask.dot(a, b)
>>> np.asarray(c.to_memmap())[0, 0]
6.0

Batched (leading dim broadcasts, only the last two axes multiply):

>>> a3 = jask.DiskArray.from_numpy(np.ones((5, 4, 6), dtype=np.float32))
>>> b3 = jask.DiskArray.from_numpy(np.ones((5, 6, 3), dtype=np.float32))
>>> jask.dot(a3, b3).shape
(5, 4, 3)

jask.materialize(x)

Bring a disk-backed array fully into memory as a real jax.Array.

Reads the whole array in one shot - only use this when x is known to fit comfortably in memory, typically at the end of a disk-backed pipeline (e.g. right before handing a small result to jax.pmap, or computing a scalar loss). Differentiable: the backward pass writes the incoming cotangent back to a disk-backed gradient.

Parameters:: x (DiskArray) – The disk-backed array to materialize. Its full contents will be loaded into memory.
Returns:: An ordinary, in-memory jax.Array with the same shape, dtype, and values as x.
Return type:: jax.Array

Examples

>>> import numpy as np
>>> import jask
>>> jask.set_memory_budget("1GB")
>>> a = jask.DiskArray.from_numpy(np.ones((4, 4), dtype=np.float32))
>>> real = jask.materialize(a)
>>> type(real).__name__
'ArrayImpl'