From 938ce7030a62b9a4acc48791f09c65736f59e223 Mon Sep 17 00:00:00 2001 From: Michael Warner Date: Tue, 8 Jul 2025 18:25:44 -0700 Subject: [PATCH 1/5] Added random operations --- conftest.py | 1 + csdl_alpha/backends/jax/graph_to_jax.py | 70 ++++++++++- csdl_alpha/src/operations/__init__.py | 4 + .../src/operations/operation_subclasses.py | 7 ++ csdl_alpha/src/operations/random/__init__.py | 0 csdl_alpha/src/operations/random/bernoulli.py | 115 ++++++++++++++++++ csdl_alpha/src/operations/random/normal.py | 65 ++++++++++ 7 files changed, 256 insertions(+), 6 deletions(-) create mode 100644 csdl_alpha/src/operations/random/__init__.py create mode 100644 csdl_alpha/src/operations/random/bernoulli.py create mode 100644 csdl_alpha/src/operations/random/normal.py diff --git a/conftest.py b/conftest.py index 1697964..6990c1c 100644 --- a/conftest.py +++ b/conftest.py @@ -17,6 +17,7 @@ def pytest_addoption(parser): additional_modules += list((package_loc / "csdl_alpha" / "src" / "operations" / "sparse").glob("*.py")) additional_modules += list((package_loc / "csdl_alpha" / "src" / "operations" / "derivatives").glob("*.py")) additional_modules += list((package_loc / "csdl_alpha" / "src" / "operations" / "special").glob("*.py")) +additional_modules += list((package_loc / "csdl_alpha" / "src" / "operations" / "random").glob("*.py")) def pytest_collect_file(file_path, path, parent): if file_path in additional_modules: diff --git a/csdl_alpha/backends/jax/graph_to_jax.py b/csdl_alpha/backends/jax/graph_to_jax.py index 69b0800..c06a884 100644 --- a/csdl_alpha/backends/jax/graph_to_jax.py +++ b/csdl_alpha/backends/jax/graph_to_jax.py @@ -1,14 +1,15 @@ +from multiprocessing import Array from ...src.graph.graph import Graph from ...src.graph.variable import Variable from ...utils.inputs import listify_variables from csdl_alpha.src.operations.loops.loop import Loop from csdl_alpha.src.operations.implicit_operations.implicit_operation import ImplicitOperation +from csdl_alpha.src.operations.operation_subclasses import RandomOperation import numpy as np from typing import Union, Callable - # Get the graph def get_jax_inputs(node, all_jax_variables:dict, return_dict = False)->list: import jax.numpy as jnp @@ -72,6 +73,9 @@ def create_jax_function( """ current_graph = graph import jax.numpy as jnp + from jax import random + from jax._src.typing import Array + KeyArray = Array # Type alias for JAX PRNGKey inputs = list(inputs) outputs = list(outputs) @@ -99,7 +103,22 @@ def create_jax_function( # all_sorted_nodes = build_derivative_node_order(current_graph, outputs, inputs, reverse=False) # Build the JAX function itself - def jax_function(*args)->list: + def jax_function(*args, prng_key:KeyArray=None)->list: + """JAX function that computes the outputs given the inputs. + + Parameters + ---------- + args : list + The input values to the function, in the same order as the inputs. + prng_key : KeyArray, optional + JAX PRNG key for random operations. If the graph contains random operations, this must be provided. + + Returns + ------- + list + The computed outputs, in the same order as the outputs. + """ + # Set the input values all_jax_variables = {} relevant_nodes = set() @@ -140,6 +159,21 @@ def jax_function(*args)->list: if fill_outputs[output_node] is None: raise ValueError(f"Jax function error with node {node.info()}: Output {output_node.info()} was not filled") all_jax_variables[output_node] = fill_outputs[output_node] + elif isinstance(node, RandomOperation): + # Random operation get a key as an argument to their compute_jax function + if prng_key is None: + raise ValueError(f"Jax function error with node {node.info()}: Random operation requires a PRNG key, but got None") + # split the key for each random operation + prng_key, node_key = random.split(prng_key, 2) + + jax_inputs = get_jax_inputs(node, all_jax_variables) + # key is always the first argument for random operations, so we prepend it to jax_inputs + jax_inputs = [node_key] + jax_inputs + jax_outputs = node.compute_jax(*jax_inputs) + if isinstance(jax_outputs, jnp.ndarray): + jax_outputs = (jax_outputs,) + update_jax_variables(node, jax_outputs, all_jax_variables) + else: jax_inputs = get_jax_inputs(node, all_jax_variables) jax_outputs = node.compute_jax(*jax_inputs) # EVERY CSDL OPERATIONS NEEDS THIS FUNCTION @@ -161,7 +195,7 @@ def create_jax_interface( graph:Graph = None, device:str='gpu', enable_f64:bool=True, - name = 'jax_interface')->Callable[[dict[Variable, np.ndarray]], dict[Variable, np.ndarray]]: + name = 'jax_interface'): """_summary_ Parameters @@ -176,10 +210,21 @@ def create_jax_interface( Returns ------- jax interface: Callable - A function with type signature: jax_interface(dict[Variable, np.array])->dict[Variable, np.array], where the input and output variables must match the inputs and outputs respectively. + A function with type signature: + jax_interface(dict[Variable, np.array], prng_key:KeyArray=None) -> dict[Variable, np.array], + where the input and output variables must match the inputs and outputs respectively. + + Notes + ----- + The returned jax_interface function accepts an optional keyword argument: + + prng_key : KeyArray, optional + JAX PRNG key for random operations. Required if the graph contains random operations. """ import jax import csdl_alpha as csdl + from jax._src.typing import Array + KeyArray = Array # Type alias for JAX PRNGKey # import os # os.environ['XLA_FLAGS'] = ( @@ -216,12 +261,25 @@ def create_jax_interface( jax_function = jax.jit(jax_function, device=device) # Create the JAX interface - def jax_interface(inputs_dict:dict[Variable, np.ndarray])->dict[Variable, np.ndarray]: + def jax_interface(inputs_dict:dict[Variable, np.ndarray], prng_key:KeyArray=None)->dict[Variable, np.ndarray]: + """ + Parameters + ---------- + inputs_dict : dict[Variable, np.ndarray] + Dictionary mapping input Variables to their values. + prng_key : KeyArray, optional + JAX PRNG key for random operations. Required if the graph contains random operations. + + Returns + ------- + outputs_dict : dict[Variable, np.ndarray] + Dictionary mapping output Variables to their computed values. + """ jax_interface_inputs = [] # print('INPUTS:') for input_var in inputs: jax_interface_inputs.append(jax.numpy.array(inputs_dict[input_var])) - jax_outputs = jax_function(*jax_interface_inputs) + jax_outputs = jax_function(*jax_interface_inputs, prng_key=prng_key) # jax_grad_outputs = jax_grad(*jax_interface_inputs) diff --git a/csdl_alpha/src/operations/__init__.py b/csdl_alpha/src/operations/__init__.py index 7baa440..5c2a02b 100644 --- a/csdl_alpha/src/operations/__init__.py +++ b/csdl_alpha/src/operations/__init__.py @@ -57,5 +57,9 @@ # Special operations from .special.bessel import bessel +# Random operations +from .random.bernoulli import bernoulli +from .random.normal import normal + # other from .subop import subop diff --git a/csdl_alpha/src/operations/operation_subclasses.py b/csdl_alpha/src/operations/operation_subclasses.py index 95538eb..715cc3e 100644 --- a/csdl_alpha/src/operations/operation_subclasses.py +++ b/csdl_alpha/src/operations/operation_subclasses.py @@ -237,6 +237,13 @@ def evaluate_composed(self, *args): cotangents.accumulate(input_var, wrt_derivs[i]) i+=1 +@set_properties() +class RandomOperation(Operation): + """ + Base class for random operations. + """ + def evaluate_vjp(self, *args): + raise NotImplementedError("Random operations do not support VJP evaluation.") # def expand_subgraph(evaluate_function): # """ diff --git a/csdl_alpha/src/operations/random/__init__.py b/csdl_alpha/src/operations/random/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/csdl_alpha/src/operations/random/bernoulli.py b/csdl_alpha/src/operations/random/bernoulli.py new file mode 100644 index 0000000..bfc44da --- /dev/null +++ b/csdl_alpha/src/operations/random/bernoulli.py @@ -0,0 +1,115 @@ +import pytest +from csdl_alpha.src.operations.operation_subclasses import RandomOperation +from csdl_alpha.src.graph.variable import Variable +from csdl_alpha.utils.inputs import validate_and_variablize +from csdl_alpha.utils.typing import VariableLike +import csdl_alpha.utils.testing_utils as csdl_tests +import numpy as np + + +class Bernoulli(RandomOperation): + def __init__(self, p:Variable, shape:tuple): + """Initialize the Bernoulli distribution. + + Parameters + ---------- + p : Variable + The probability of success (i.e., the probability of getting 1). + shape : tuple + The shape of the output array. + """ + super().__init__(p) + self.shape = shape + self.name = 'bernoulli' + self.set_dense_outputs((self.shape,)) + + def compute_inline(self, p): + return np.random.binomial(1, p, self.shape) + + def compute_jax(self, key, p): + """Compute the Bernoulli distribution using JAX. + + Parameters + ---------- + key : jax.random.PRNGKey + The random key for JAX. + + Returns + ------- + jax.numpy.ndarray + An array of shape `self.shape` with values drawn from a Bernoulli distribution. + """ + from jax import random + + return random.bernoulli(key, p, shape=self.shape) + + +def bernoulli(p:VariableLike, shape:tuple) -> Variable: + """Create a Bernoulli distribution variable. + + Parameters + ---------- + p : VariableLike + The probability of success (i.e., the probability of getting 1). + shape : tuple + The shape of the output array. + + Returns + ------- + Variable + A variable representing the Bernoulli distribution, with shape `shape`. + """ + p = validate_and_variablize(p) + return Bernoulli(p, shape).finalize_and_return_outputs() + + + +class TestBernoulli(csdl_tests.CSDLTest): + def test_functionality(self): + self.prep(always_build_inline=True) + + import csdl_alpha as csdl + import numpy as np + + p_val = 0.7 + shape = (3, 4) + + p = csdl.Variable(value=p_val) + bernoulli_var = bernoulli(p, shape) + + # Check the shape of the result + assert bernoulli_var.shape == shape, f"Expected shape {shape}, got {bernoulli_var.shape}" + + # Check that the values are either 0 or 1 + assert np.all(np.isin(bernoulli_var.value, [0, 1])), f"Expected values to be 0 or 1, got {bernoulli_var.value}" + + def test_jax_interface(self): + self.prep() + + import csdl_alpha as csdl + import jax.numpy as jnp + from jax import random + + p_val = 0.7 + shape = (3, 4) + + p = csdl.Variable(value=p_val) + bernoulli_var = bernoulli(p, shape) + + interface = csdl.jax.create_jax_interface(p, bernoulli_var) + outputs = interface({p:p.value}, prng_key=random.PRNGKey(42)) + + def test_derivative_error(self): + self.prep() + + import csdl_alpha as csdl + + p_val = 0.7 + shape = (3, 4) + + p = csdl.Variable(value=p_val) + bernoulli_var = bernoulli(p, shape) + + # Check that the derivative raises an error + with pytest.raises(ValueError): + csdl.derivative(bernoulli_var, p) \ No newline at end of file diff --git a/csdl_alpha/src/operations/random/normal.py b/csdl_alpha/src/operations/random/normal.py new file mode 100644 index 0000000..f063f67 --- /dev/null +++ b/csdl_alpha/src/operations/random/normal.py @@ -0,0 +1,65 @@ +import pytest +from csdl_alpha.src.operations.operation_subclasses import RandomOperation +from csdl_alpha.src.graph.variable import Variable +from csdl_alpha.utils.inputs import validate_and_variablize +from csdl_alpha.utils.typing import VariableLike +import csdl_alpha.utils.testing_utils as csdl_tests +import numpy as np + +class Normal(RandomOperation): + def __init__(self, shape:tuple): + super().__init__() + self.shape = shape + self.name = 'normal' + self.set_dense_outputs((self.shape,)) + + def compute_inline(self): + return np.random.normal(size=self.shape) + + def compute_jax(self, key): + """Compute the normal distribution using JAX. + + Parameters + ---------- + key : jax.random.PRNGKey + The random key for JAX. + + Returns + ------- + jax.numpy.ndarray + An array of shape `self.shape` with values drawn from a normal distribution. + """ + from jax import random + return random.normal(key, shape=self.shape) + +def normal(shape:tuple) -> Variable: + """Create a normal distribution variable. + + Parameters + ---------- + shape : tuple + The shape of the output array. + + Returns + ------- + Variable + A variable representing the normal distribution, with shape `shape`. + """ + return Normal(shape).finalize_and_return_outputs() + + +class TestNormal(csdl_tests.CSDLTest): + def test_functionality(self): + self.prep(always_build_inline=True) + + import csdl_alpha as csdl + import numpy as np + + shape = (3, 4) + normal_var = normal(shape) + + # Check the shape of the result + assert normal_var.shape == shape, f"Expected shape {shape}, got {normal_var.shape}" + + # Check that the values are normally distributed + assert np.all(np.isfinite(normal_var.value)), "Expected all values to be finite" From 478ffb8ab3350d54d2f4cf0251c26980c6534a95 Mon Sep 17 00:00:00 2001 From: mwarner Date: Sat, 18 Jul 2026 09:35:46 -0700 Subject: [PATCH 2/5] added vmap_method="sequential" to pure callbacks --- csdl_alpha/backends/jax/utils.py | 3 ++- csdl_alpha/src/operations/custom/custom.py | 6 ++++-- csdl_alpha/src/operations/subop.py | 3 ++- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/csdl_alpha/backends/jax/utils.py b/csdl_alpha/backends/jax/utils.py index 12e3ef1..eceb8a5 100644 --- a/csdl_alpha/backends/jax/utils.py +++ b/csdl_alpha/backends/jax/utils.py @@ -15,7 +15,8 @@ def new_inline_func(*args_in): output = jax.pure_callback( new_inline_func, [jax.ShapeDtypeStruct(output.shape, np.float64) for output in operation.outputs], - *args) + *args, + vmap_method="sequential") return tuple(output) diff --git a/csdl_alpha/src/operations/custom/custom.py b/csdl_alpha/src/operations/custom/custom.py index 63e7d03..ccd579a 100644 --- a/csdl_alpha/src/operations/custom/custom.py +++ b/csdl_alpha/src/operations/custom/custom.py @@ -97,7 +97,8 @@ def new_inline_func(*args): output = jax.pure_callback( new_inline_func, [jax.ShapeDtypeStruct(self.output_dict[output_var].shape, dtype) for output_var in self.output_dict], - *args) + *args, + vmap_method="sequential") # if len(output) == 1: # output = output[0] return tuple(output) @@ -463,7 +464,8 @@ def new_inline_func(*args): output = jax.pure_callback( new_inline_func, [jax.ShapeDtypeStruct(in_cot.shape, dtype) for in_cot in self.input_cotangents], - *args) + *args, + vmap_method="sequential") # if len(output) == 1: # output = output[0] diff --git a/csdl_alpha/src/operations/subop.py b/csdl_alpha/src/operations/subop.py index cf923ef..b01a614 100644 --- a/csdl_alpha/src/operations/subop.py +++ b/csdl_alpha/src/operations/subop.py @@ -41,7 +41,8 @@ def compute_jax(self, *args): if self.jit: output = jax.pure_callback(jax.jit(jax_function), [jax.ShapeDtypeStruct(output.shape, np.float64) for output in self.outputs], - *args) + *args, + vmap_method="sequential") return tuple(output) # return tuple(jax.jit(jax_function)(*args)) else: From 67cbc62832586b75dcb2fec5d97c39a9befb88dc Mon Sep 17 00:00:00 2001 From: mwarner Date: Thu, 27 Aug 2026 16:03:16 -0700 Subject: [PATCH 3/5] numpy 2 compat: reshape size-1 RHS to slice target in SetVarIndex.compute_inline NumPy 2 no longer silently squeezes a size-1 array assigned into a scalar slot (`x[i] = array([v])` raises "setting an array element with a sequence"). This broke the reverse-mode derivative seed (`seed.set(csdl.slice[i], 1.0)`) and hence all of compute_totals / check_totals under numpy>=2. Co-Authored-By: Claude Sonnet 5 --- csdl_alpha/src/operations/set_get/setindex.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/csdl_alpha/src/operations/set_get/setindex.py b/csdl_alpha/src/operations/set_get/setindex.py index f375fd1..a145c50 100644 --- a/csdl_alpha/src/operations/set_get/setindex.py +++ b/csdl_alpha/src/operations/set_get/setindex.py @@ -31,7 +31,13 @@ def __init__( def compute_inline(self, x, y, *slice_args): x_updated = x.copy() - x_updated[self.slice.evaluate(*slice_args)] = y + eval_slice = self.slice.evaluate(*slice_args) + # NumPy 2 no longer silently squeezes a size-1 RHS into a scalar target; + # reshape y to the slice target when the total size matches. + tgt_shape = np.shape(x_updated[eval_slice]) + if np.shape(y) != tgt_shape and np.size(y) == int(np.prod(tgt_shape)): + y = np.asarray(y).reshape(tgt_shape) + x_updated[eval_slice] = y return x_updated # # Set item could add over duplicate indices. From b2b5b9405706c2ab066cc8936bb5573ba389ecd0 Mon Sep 17 00:00:00 2001 From: Michael Warner Date: Tue, 15 Sep 2026 14:18:41 -0700 Subject: [PATCH 4/5] Fix PRNG key not forwarded to cost-analysis retrace in jax_interface jax_interface() computed outputs correctly by passing prng_key to jax_function, but the debug cost-analysis block re-traced jax_function without it, causing RandomOperation nodes to raise "Random operation requires a PRNG key, but got None" for any graph containing e.g. bernoulli or normal. Co-Authored-By: Claude Sonnet 5 --- csdl_alpha/backends/jax/graph_to_jax.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/csdl_alpha/backends/jax/graph_to_jax.py b/csdl_alpha/backends/jax/graph_to_jax.py index e6fe239..42d7d62 100644 --- a/csdl_alpha/backends/jax/graph_to_jax.py +++ b/csdl_alpha/backends/jax/graph_to_jax.py @@ -283,7 +283,7 @@ def jax_interface(inputs_dict:dict[Variable, np.ndarray], prng_key:KeyArray=None #### Potential analysis tools #### ## ----- compiled func cost estimates ----- - traced = jax_function.trace(*jax_interface_inputs) + traced = jax_function.trace(*jax_interface_inputs, prng_key=prng_key) lowered = traced.lower() compiled = lowered.compile() # for compiled_costs in compiled.cost_analysis(): From 78887bd696e79f64eacfdbbaaf15b5c49ea83901 Mon Sep 17 00:00:00 2001 From: Michael Warner Date: Tue, 15 Sep 2026 14:51:01 -0700 Subject: [PATCH 5/5] Tighten numpy 2 squeeze-compat guard in SetVarIndex.compute_inline The prior guard reshaped y whenever its total size matched the slice target's, regardless of shape. That's broader than intended: a genuinely mismatched-but-same-size y (e.g. a transposed array from a bug elsewhere) would be silently reshaped instead of raising numpy's broadcast error. Gate on the target also being size-1, matching the actual scenario being worked around (a size-1 RHS squeezed into a size-1 slot), so any other mismatch still fails loudly. Co-Authored-By: Claude Sonnet 5 --- csdl_alpha/src/operations/set_get/setindex.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/csdl_alpha/src/operations/set_get/setindex.py b/csdl_alpha/src/operations/set_get/setindex.py index a145c50..242d330 100644 --- a/csdl_alpha/src/operations/set_get/setindex.py +++ b/csdl_alpha/src/operations/set_get/setindex.py @@ -32,10 +32,12 @@ def __init__( def compute_inline(self, x, y, *slice_args): x_updated = x.copy() eval_slice = self.slice.evaluate(*slice_args) - # NumPy 2 no longer silently squeezes a size-1 RHS into a scalar target; - # reshape y to the slice target when the total size matches. tgt_shape = np.shape(x_updated[eval_slice]) - if np.shape(y) != tgt_shape and np.size(y) == int(np.prod(tgt_shape)): + # NumPy 2 no longer silently squeezes a size-1 RHS into a size-1 (e.g. scalar) + # target; reshape y to match so the assignment succeeds. Gated on the target + # also being size-1 so a genuinely mismatched y (a real bug elsewhere) still + # raises numpy's broadcast error instead of being silently reshaped. + if np.size(y) == 1 and int(np.prod(tgt_shape)) == 1 and np.shape(y) != tgt_shape: y = np.asarray(y).reshape(tgt_shape) x_updated[eval_slice] = y return x_updated