Tensor Abstractions: NumPy, PyTorch, JAX

Overview

Tensor libraries define different abstractions for multi-dimensional arrays: - NumPy: CPU, mutable, static dtype - PyTorch: CPU/GPU, mutable, autograd - JAX: GPU/TPU, immutable, functional, traced - CuPy: GPU variant of NumPy - TensorFlow: Different semantics from PyTorch

Understanding these differences is critical for building ML systems that work across frameworks.

---

The Tensor Protocol: Array-Like Objects

NumPy Array Interface

All major frameworks implement the NumPy array interface:

import numpy as np
import torch
import jax.numpy as jnp

# All implement __array__ protocol
x_np = np.array([1, 2, 3])
x_torch = torch.tensor([1, 2, 3])
x_jax = jnp.array([1, 2, 3])

# Convert to NumPy
print(np.asarray(x_torch))  # Calls x_torch.__array__()
print(np.asarray(x_jax))    # Calls x_jax.__array__()

__array_interface__ Protocol

Low-level protocol for zero-copy sharing:

class MyTensor:
    def __init__(self, data):
        self.data = np.array(data)

    @property
    def __array_interface__(self):
        """Enable zero-copy access to underlying buffer."""
        return {
            'shape': self.data.shape,
            'typestr': self.data.dtype.str,
            'data': (self.data.data, False),
            'version': 3,
        }

# Other libraries can access data without copying
t = MyTensor([1, 2, 3])
arr = np.asarray(t)  # Zero-copy if possible

__array_ufunc__ Protocol

Allows frameworks to override NumPy operations:

class CustomTensor:
    def __init__(self, data):
        self.data = np.array(data)

    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
        """Called when NumPy ufunc operates on this object."""
        if ufunc == np.add:
            # Custom addition logic
            return CustomTensor(self.data + inputs[1].data)
        return NotImplemented

# Usage
t1 = CustomTensor([1, 2])
t2 = CustomTensor([3, 4])
result = np.add(t1, t2)  # Calls t1.__array_ufunc__

---

PyTorch Tensors

Fundamental Characteristics

import torch

# Creation
x = torch.tensor([1, 2, 3])           # From Python list
x = torch.zeros(3, 4)                 # Zeros
x = torch.randn(3, 4)                 # Random normal
x = torch.arange(10)                  # Range
x = torch.linspace(0, 1, 100)         # Linspace

# Device placement
x = x.to('cuda')                      # Move to GPU
x = x.to(torch.device('cuda:0'))     # Explicit device
x = x.cuda()                          # Shorthand
x = x.cpu()                           # Move to CPU

# Data type
x = x.float()     # torch.float32
x = x.double()    # torch.float64
x = x.half()      # torch.float16
x = x.int()       # torch.int32
x = torch.tensor([1, 2], dtype=torch.float32)

Mutability and Gradients

# PyTorch tensors are mutable
x = torch.tensor([1.0, 2.0, 3.0])
x[0] = 5.0  # In-place modification
print(x)    # tensor([5., 2., 3.])

# Requires grad for autograd
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x ** 2
loss = y.sum()
loss.backward()
print(x.grad)  # tensor([2., 4.])

# In-place ops affect gradient computation
x = torch.tensor([1.0, 2.0], requires_grad=True)
x += 1  # In-place add
# Later operations will see modified x

Memory Layout

# Row-major (C-contiguous) vs Column-major (Fortran-contiguous)
x = torch.randn(3, 4)
print(x.is_contiguous())  # True (C-contiguous by default)

# Transposed is not contiguous
y = x.T
print(y.is_contiguous())  # False

# Some operations require contiguous
z = y.contiguous()  # Makes copy if needed
print(z.is_contiguous())  # True

# Important for GPU kernels and backward pass

Indexing and Slicing

x = torch.randn(3, 4, 5)

# Basic indexing
x[0]           # First row
x[0, 1]        # Element [0, 1]
x[:, 1, :]     # Middle column, all others

# Advanced indexing
idx = torch.tensor([0, 2])
x[idx]         # Rows 0 and 2

# Fancy indexing
mask = x > 0
x[mask]        # Elements > 0 (flattened)

# Indexing creates views (memory shared)
y = x[:, 0]
y[0] = 999     # Modifies x too!

---

JAX Arrays

Immutability

JAX arrays are immutable by design (enable JIT compilation):

import jax.numpy as jnp

x = jnp.array([1, 2, 3])

# Can't modify in-place
# x[0] = 5  # TypeError!

# Must use functional updates
x = x.at[0].set(5)
print(x)  # [5, 2, 3] (new array)

# Batch operations
x = x.at[jnp.array([0, 2])].set(99)
print(x)  # [99, 2, 99]

# Add instead of increment
x = x.at[0].add(10)
print(x)  # [109, 2, 99]

Shape Tracing

JAX uses shape tracing during JIT compilation:

import jax
import jax.numpy as jnp

@jax.jit
def sum_array(x):
    return jnp.sum(x)

# Shape is traced at compile time
x = jnp.array([1, 2, 3])
result = sum_array(x)  # JIT compiles with shape (3,)

# Different shape requires recompilation
y = jnp.array([1, 2, 3, 4, 5])
result = sum_array(y)  # Recompiles with shape (5,)

Functional Transformations

import jax
import jax.numpy as jnp

def loss_fn(params, x, y):
    pred = jnp.dot(x, params)
    return jnp.mean((pred - y) ** 2)

# Gradient: function → function (same interface)
grad_fn = jax.grad(loss_fn)
grads = grad_fn(params, x, y)  # Returns gradients w.r.t. params

# Vectorize: function → batched function
batched_loss = jax.vmap(loss_fn, in_axes=(None, 0, 0))
losses = batched_loss(params, x_batch, y_batch)  # Vectorized

# Compose transformations
@jax.jit
@jax.vmap
def batched_forward(params, batch_x):
    return jnp.dot(batch_x, params)

---

NumPy Arrays

Core Characteristics

import numpy as np

# Creation
x = np.array([1, 2, 3])
x = np.zeros((3, 4))
x = np.random.randn(3, 4)

# CPU only
x = x  # No device placement

# Mutable
x[0] = 5  # In-place modification

# No autograd
# No backward() method

# Broadcasting
x = np.array([1, 2, 3](/1,-2,-3/))  # Shape (1, 3)
y = np.array([[1], [2], [3]])  # Shape (3, 1)
result = x + y  # Broadcasts to (3, 3)

Broadcasting Rules

Critical for ML code:

import numpy as np

# Rule 1: Align dimensions from right
# (3, 1, 4) broadcasts with (1, 3, 4)
# Result: (3, 3, 4)

x = np.zeros((3, 1, 4))
y = np.zeros((1, 3, 4))
result = x + y
print(result.shape)  # (3, 3, 4)

# Rule 2: Add dimensions as needed
# (3, 4) broadcasts with (4,)
x = np.zeros((3, 4))
y = np.zeros(4)
result = x + y  # y broadcasted to (1, 4) then (3, 4)
print(result.shape)  # (3, 4)

# Practical example: batch operations
batch_size = 32
seq_len = 10
vocab_size = 50000

logits = np.random.randn(batch_size, seq_len, vocab_size)
temperature = 0.5  # Scalar

scaled_logits = logits / temperature  # Broadcasting works
print(scaled_logits.shape)  # (32, 10, 50000)

---

Practical: Cross-Framework Code

Write Once, Run Anywhere

from typing import Protocol, Union
import numpy as np
import torch
import jax.numpy as jnp

class ArrayLike(Protocol):
    """Any array-like object."""
    shape: tuple
    dtype: object

    def __array__(self) -> np.ndarray:
        ...

def normalize(arr: ArrayLike) -> ArrayLike:
    """Works with NumPy, PyTorch, JAX arrays."""
    mean = np.mean(arr)
    std = np.std(arr)
    return (arr - mean) / std

# All work!
normalize(np.array([1, 2, 3]))
normalize(torch.tensor([1, 2, 3]))
normalize(jnp.array([1, 2, 3]))

Framework-Agnostic Training Loop

import numpy as np
import torch

def train_step(model, x, y, optimizer, framework='torch'):
    """Training step works with PyTorch or NumPy models."""

    if framework == 'torch':
        # PyTorch path
        optimizer.zero_grad()
        pred = model(x)
        loss = torch.nn.functional.mse_loss(pred, y)
        loss.backward()
        optimizer.step()
        return float(loss)

    elif framework == 'numpy':
        # NumPy path (manual backprop)
        pred = model.forward(x)
        loss = np.mean((pred - y) ** 2)
        grad = 2 * (pred - y) / len(y)
        model.backward(grad)
        model.update()
        return float(loss)

Conversion Utilities

import numpy as np
import torch
import jax.numpy as jnp

def to_numpy(x):
    """Convert any array to NumPy."""
    if isinstance(x, np.ndarray):
        return x
    elif isinstance(x, torch.Tensor):
        return x.detach().cpu().numpy()
    elif isinstance(x, jnp.ndarray):
        return np.array(x)
    else:
        return np.asarray(x)

def to_torch(x, device='cpu'):
    """Convert any array to PyTorch."""
    if isinstance(x, torch.Tensor):
        return x.to(device)
    else:
        x_np = to_numpy(x)
        return torch.from_numpy(x_np).to(device)

def to_jax(x):
    """Convert any array to JAX."""
    if isinstance(x, jnp.ndarray):
        return x
    else:
        x_np = to_numpy(x)
        return jnp.array(x_np)

# Usage
x_np = np.array([1, 2, 3])
x_torch = to_torch(x_np, device='cuda')
x_jax = to_jax(x_torch)
x_np = to_numpy(x_jax)

---

Tensor Shape and Dimension Operations

Common Operations

import torch

x = torch.randn(3, 4, 5)

# Shape inspection
print(x.shape)       # torch.Size([3, 4, 5])
print(x.ndim)        # 3
print(x.size())      # torch.Size([3, 4, 5])

# Reshape
y = x.reshape(12, 5)     # (3, 4, 5) → (12, 5)
y = x.view(12, 5)        # Requires contiguous (faster)
y = x.flatten()          # Flatten to 1D

# Transpose
y = x.T                  # Full transpose
y = x.transpose(0, 2)   # Swap dims 0 and 2

# Squeeze/Unsqueeze
x_sq = torch.randn(1, 3, 1, 4)
y = x_sq.squeeze()      # Remove dims of size 1: (3, 4)
y = x.unsqueeze(0)      # Add dimension at 0: (1, 3, 4, 5)

# Permute
y = x.permute(2, 0, 1)  # Reorder: (5, 3, 4)

# Expand
y = x.expand(3, 4, 10)  # Only works if target size ≥ original

# Repeat
y = x.repeat(2, 1, 1)   # Copy tensor: (6, 4, 5)

Broadcasting vs Reshaping

import torch

# Broadcasting: create virtual copies (no memory cost)
x = torch.zeros(3, 1, 4)
y = torch.zeros(1, 3, 4)
result = x + y  # Broadcasts to (3, 3, 4) without copying

# Reshaping: change layout (may copy)
z = torch.randn(12)
w = z.reshape(3, 4)  # May or may not copy depending on layout

# Difference in inference:
# - Broadcasting: efficient for batch operations
# - Reshaping: needed to change data layout

---

Device Management Patterns

Automatic Device Handling

import torch

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 5)

    @property
    def device(self):
        """Infer device from parameters."""
        return next(self.parameters()).device

    def to_device(self, x):
        """Move input to same device as model."""
        return x.to(self.device)

    def forward(self, x):
        x = self.to_device(x)
        return self.linear(x)

# Usage
model = Model()
model.to('cuda')

# Forward works even if input is on CPU
x_cpu = torch.randn(3, 10)
output = model(x_cpu)  # Automatically moved to GPU

Mixed Precision with Automatic Device Placement

import torch

# Model on GPU, but some computations in float32
model = torch.nn.Linear(10, 5).to('cuda').half()

# Input automatically cast to float16
x = torch.randn(3, 10, device='cuda', dtype=torch.float16)
y = model(x)  # Computed in float16

# Some operations require float32
with torch.autocast(device_type='cuda', dtype=torch.float32):
    # All operations in float32
    z = torch.nn.functional.softmax(y, dim=-1)

---

Practical: Implementing Tensor-Agnostic Code

from typing import Union, Optional
import numpy as np
import torch

Array = Union[np.ndarray, torch.Tensor]

class TensorProcessor:
    """Works with NumPy or PyTorch tensors."""

    @staticmethod
    def is_torch(x):
        return isinstance(x, torch.Tensor)

    @staticmethod
    def zero_like(x: Array) -> Array:
        """Create zero tensor same shape/dtype as input."""
        if TensorProcessor.is_torch(x):
            return torch.zeros_like(x)
        else:
            return np.zeros_like(x)

    @staticmethod
    def concatenate(arrays: list[Array], axis: int = 0) -> Array:
        """Concatenate along axis."""
        if TensorProcessor.is_torch(arrays[0]):
            return torch.cat(arrays, dim=axis)
        else:
            return np.concatenate(arrays, axis=axis)

    @staticmethod
    def normalize(x: Array) -> Array:
        """Normalize to [0, 1]."""
        x_min = x.min()
        x_max = x.max()
        return (x - x_min) / (x_max - x_min + 1e-8)

# Usage
x_np = np.array([1, 2, 3])
x_torch = torch.tensor([1, 2, 3])

norm_np = TensorProcessor.normalize(x_np)
norm_torch = TensorProcessor.normalize(x_torch)

---

Summary

Feature	NumPy	PyTorch	JAX
Mutability	Mutable	Mutable	Immutable
Devices	CPU only	CPU/GPU/TPU	GPU/TPU
Autograd	No	Yes	Via transformations
JIT	No	TorchScript	Native
Broadcasting	Yes	Yes	Yes
Functional	No	No	Yes
Array Interface	Yes	Yes	Yes
Use Case	General	Production ML	Research/JAX

---

Related Topics

• 02 Autograd Implementation - How PyTorch tracks gradients
• 05 Custom Operators - Extending tensors
• 06 Inference Optimization Patterns - Efficient tensor operations
• 00 Readme - Tensor implementation in C/C++