Add latest Python/JAX/NumPy CI job; bump version to 0.1.1

.github/workflows/actions.yml

@@ -42,3 +42,27 @@ jobs:
     - name: Test with pytest
       run: |
         pytest
+
+  # Canary run on the newest stable Python with the newest JAX and NumPy, ignoring the pins in requirements.txt
+  latest:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up latest Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.x"
+        check-latest: true
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -U jax numpy
+        pip install -e .
+        python -c "import sys, jax, numpy; print('python', sys.version); print('jax', jax.__version__); print('numpy', numpy.__version__)"
+
+    - name: Test with pytest
+      run: |
+        pytest

csdl_alpha/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.1.0'
+__version__ = '0.1.1'
 
 
 # Standard operations

csdl_alpha/src/graph/variable.py

@@ -321,7 +321,7 @@ def __getitem__(self, indices:Union[Slice, tuple[list[int], int, slice]]) -> 'Va
         array([[3., 4., 5.]])
         >>> x[:].shape
         (2, 3)
-        >>> np.all(x[:].value == x.value)
+        >>> np.array_equal(x[:].value, x.value)
         True
         >>> x[1:2,:-1].shape
         (1, 2)

csdl_alpha/src/operations/custom/tests/test_custom_implicit.py

@@ -54,7 +54,7 @@ def compute_jacvec_product(self, inputs, outputs, d_inputs, d_outputs, d_residua
         # for mode = rev
         # d_residuals --> d_inputs
         if mode == 'rev':
-            d_inputs['a'][1,1] = d_residuals['x']*x_solved**2 + d_residuals['y']*y_solved**2
+            d_inputs['a'][1,1] = np.sum(d_residuals['x']*x_solved**2 + d_residuals['y']*y_solved**2)
             d_inputs['b'] = d_residuals['x']*x_solved + d_residuals['y']*(y_solved+y_solved**2)
             d_inputs['c'] = d_residuals['x'] + d_residuals['y']/2.0
 

csdl_alpha/src/variable_group.py

@@ -1,4 +1,5 @@
 from typing import Union
+from dataclasses import fields, is_dataclass
 import warnings
 from csdl_alpha.src.graph.variable import Variable
 from csdl_alpha.utils.inputs import variablize
@@ -74,7 +75,7 @@ def add_check(self, name:str, type=None, shape:tuple=None, variablize:bool=False
         ValueError
             If parameters for the variable with the given name are already declared.
         """
-        if not name in self.__annotations__:
+        if not is_dataclass(self) or name not in {field.name for field in fields(self)}:
             raise ValueError(f"Variable {name} not found in the group.")
         if name in self._metadata:
             raise ValueError(f"Checks for variable {name} already declared.")

recipe/recipe.yaml

@@ -2,7 +2,7 @@ schema_version: 1
 
 package:
   name: csdl_alpha
-  version: "0.0.0a2"
+  version: "0.1.1"
 
 source:
   path: ..