Derivative of JacobianInverse fails

[michalhabera]

import numpy as np

import ufl
import dolfinx.fem
from dolfinx.mesh import create_mesh
from mpi4py import MPI


domain = ufl.Mesh(ufl.VectorElement("Lagrange", ufl.Cell("interval", 3), 2))
mesh = create_mesh(MPI.COMM_WORLD, np.array([[0, 2, 1], [2, 4, 3]]), np.array(
    [[0.0, 0.0, 0.0], [0.25, 0.0, 0.0], [0.5, 0.0, 0.0], [0.75, 0.0, 0.0], [1.0, 0.0, 0.0]]), domain)

B = dolfinx.fem.TensorFunctionSpace(mesh, ("DG", 2), (mesh.topology.dim, mesh.topology.dim))
K = ufl.geometry.JacobianInverse(mesh).T
S = dolfinx.fem.form(ufl.inner(ufl.grad(K), ufl.grad(K)) * ufl.dx(mesh))

s = dolfinx.fem.assemble_scalar(S)

fails with

...
  File "/usr/local/lib/python3.9/dist-packages/dolfinx/jit.py", line 209, in ffcx_jit
    r = ffcx.codegeneration.jit.compile_forms([ufl_object], parameters=p_ffcx, **p_jit)
  File "/home/ffcx/ffcx/codegeneration/jit.py", line 168, in compile_forms
    impl = _compile_objects(decl, forms, form_names, module_name, p, cache_dir,
  File "/home/ffcx/ffcx/codegeneration/jit.py", line 232, in _compile_objects
    _, code_body = ffcx.compiler.compile_ufl_objects(ufl_objects, prefix=module_name, parameters=parameters)
  File "/home/ffcx/ffcx/compiler.py", line 103, in compile_ufl_objects
    ir = compute_ir(analysis, object_names, prefix, parameters, visualise)
  File "/home/ffcx/ffcx/ir/representation.py", line 94, in compute_ir
    irs = [
  File "/home/ffcx/ffcx/ir/representation.py", line 95, in <listcomp>
    _compute_integral_ir(fd, i, analysis.element_numbers, integral_names, finite_element_names,
  File "/home/ffcx/ffcx/ir/representation.py", line 360, in _compute_integral_ir
    integral_ir = compute_integral_ir(itg_data.domain.ufl_cell(), itg_data.integral_type,
  File "/home/ffcx/ffcx/ir/integral.py", line 84, in compute_integral_ir
    mt_table_reference = build_optimized_tables(
  File "/home/ffcx/ffcx/ir/elementtables.py", line 346, in build_optimized_tables
    t = get_ffcx_table_values(quadrature_rule.points, cell,
  File "/home/ffcx/ffcx/ir/elementtables.py", line 108, in get_ffcx_table_values
    tbl = tbl[basix_index(*derivative_counts)]
IndexError: index 4 is out of bounds for axis 0 with size 3

@mscroggs Any idea why most likely derivative_counts is wrong for this (a bit exotic) ufl symbol?