fix flake8

kinnala · Jan 23, 2024 · 3b0c9db · 3b0c9db
1 parent 85ecd81
commit 3b0c9db
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/skfem/experimental/autodiff/__init__.py b/skfem/experimental/autodiff/__init__.py
@@ -114,9 +114,12 @@ def assemble(self, basis, x=None, **kwargs):
         dx = basis.dx
 
         defaults = basis.default_parameters()
+        # turn defaults into JaxDiscreteField to avoid np.ndarray
+        # to jnp.ndarray promotion issues
         w = FormExtraParams({
             **{
-                k: JaxDiscreteField(*tuple(jnp.asarray(x) for x in defaults[k].astuple))
+                k: JaxDiscreteField(*tuple(jnp.asarray(x)
+                                           for x in defaults[k].astuple))
                 for k in defaults
             },
             **self._normalize_asm_kwargs(kwargs, basis),