Nutpie fails to compile

[bpajusco]

I'm trying to use nutpie to speed up an ICM GP model. I've run the model through pymc without errors but once I call the command nutpie.compile_pymc_model(icm) I get the following error message:

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function numba_funcify_Elemwise.<locals>.elemwise at 0x18dfcdea0>) found for signature:
 
 >>> elemwise(readonly array(float64, 0d, C), readonly array(float64, 0d, C), float64, array(float64, 0d, C))

For completeness here is full output

TypingError                               Traceback (most recent call last)
File <timed exec>:2

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/nutpie/compile_pymc.py:121, in compile_pymc_model(model, **kwargs)
    116 user_data = make_user_data(logp_fn_pt, shared_data)
    118 logp_numba_raw, c_sig = _make_c_logp_func(
    119     n_dim, logp_fn, user_data, shared_logp, shared_data
    120 )
--> 121 logp_numba = numba.cfunc(c_sig, **kwargs)(logp_numba_raw)
    123 def expand_draw(x, seed, chain, draw, *, shared_data):
    124     return expand_fn(x, **{name: shared_data[name] for name in shared_expand})[0]

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/decorators.py:282, in cfunc.<locals>.wrapper(func)
    280 if cache:
    281     res.enable_caching()
--> 282 res.compile()
    283 return res

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/compiler_lock.py:35, in _CompilerLock.__call__.<locals>._acquire_compile_lock(*args, **kwargs)
     32 @functools.wraps(func)
     33 def _acquire_compile_lock(*args, **kwargs):
     34     with self:
---> 35         return func(*args, **kwargs)

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/ccallback.py:67, in CFunc.compile(self)
     64 cres = self._cache.load_overload(self._sig,
     65                                  self._targetdescr.target_context)
     66 if cres is None:
---> 67     cres = self._compile_uncached()
     68     self._cache.save_overload(self._sig, cres)
     69 else:

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/ccallback.py:81, in CFunc._compile_uncached(self)
     78 sig = self._sig
     80 # Compile native function as well as cfunc wrapper
---> 81 return self._compiler.compile(sig.args, sig.return_type)

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/dispatcher.py:129, in _FunctionCompiler.compile(self, args, return_type)
    127     return retval
    128 else:
--> 129     raise retval

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/dispatcher.py:139, in _FunctionCompiler._compile_cached(self, args, return_type)
    136     pass
    138 try:
--> 139     retval = self._compile_core(args, return_type)
    140 except errors.TypingError as e:
    141     self._failed_cache[key] = e

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/dispatcher.py:152, in _FunctionCompiler._compile_core(self, args, return_type)
    149 flags = self._customize_flags(flags)
    151 impl = self._get_implementation(args, {})
--> 152 cres = compiler.compile_extra(self.targetdescr.typing_context,
    153                               self.targetdescr.target_context,
    154                               impl,
    155                               args=args, return_type=return_type,
    156                               flags=flags, locals=self.locals,
    157                               pipeline_class=self.pipeline_class)
    158 # Check typing error if object mode is used
    159 if cres.typing_error is not None and not flags.enable_pyobject:

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/compiler.py:716, in compile_extra(typingctx, targetctx, func, args, return_type, flags, locals, library, pipeline_class)
    692 """Compiler entry point
    693 
    694 Parameter
   (...)
    712     compiler pipeline
    713 """
    714 pipeline = pipeline_class(typingctx, targetctx, library,
    715                           args, return_type, flags, locals)
--> 716 return pipeline.compile_extra(func)

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/compiler.py:452, in CompilerBase.compile_extra(self, func)
    450 self.state.lifted = ()
    451 self.state.lifted_from = None
--> 452 return self._compile_bytecode()

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/compiler.py:520, in CompilerBase._compile_bytecode(self)
    516 """
    517 Populate and run pipeline for bytecode input
    518 """
    519 assert self.state.func_ir is None
--> 520 return self._compile_core()

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/compiler.py:499, in CompilerBase._compile_core(self)
    497         self.state.status.fail_reason = e
    498         if is_final_pipeline:
--> 499             raise e
    500 else:
    501     raise CompilerError("All available pipelines exhausted")

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/compiler.py:486, in CompilerBase._compile_core(self)
    484 res = None
    485 try:
--> 486     pm.run(self.state)
    487     if self.state.cr is not None:
    488         break

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/compiler_machinery.py:368, in PassManager.run(self, state)
    365 msg = "Failed in %s mode pipeline (step: %s)" % \
    366     (self.pipeline_name, pass_desc)
    367 patched_exception = self._patch_error(msg, e)
--> 368 raise patched_exception

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/compiler_machinery.py:356, in PassManager.run(self, state)
    354 pass_inst = _pass_registry.get(pss).pass_inst
    355 if isinstance(pass_inst, CompilerPass):
--> 356     self._runPass(idx, pass_inst, state)
    357 else:
    358     raise BaseException("Legacy pass in use")

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/compiler_lock.py:35, in _CompilerLock.__call__.<locals>._acquire_compile_lock(*args, **kwargs)
     32 @functools.wraps(func)
     33 def _acquire_compile_lock(*args, **kwargs):
     34     with self:
---> 35         return func(*args, **kwargs)

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/compiler_machinery.py:311, in PassManager._runPass(self, index, pss, internal_state)
    309     mutated |= check(pss.run_initialization, internal_state)
    310 with SimpleTimer() as pass_time:
--> 311     mutated |= check(pss.run_pass, internal_state)
    312 with SimpleTimer() as finalize_time:
    313     mutated |= check(pss.run_finalizer, internal_state)

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/compiler_machinery.py:273, in PassManager._runPass.<locals>.check(func, compiler_state)
    272 def check(func, compiler_state):
--> 273     mangled = func(compiler_state)
    274     if mangled not in (True, False):
    275         msg = ("CompilerPass implementations should return True/False. "
    276                "CompilerPass with name '%s' did not.")

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/typed_passes.py:105, in BaseTypeInference.run_pass(self, state)
     99 """
    100 Type inference and legalization
    101 """
    102 with fallback_context(state, 'Function "%s" failed type inference'
    103                       % (state.func_id.func_name,)):
    104     # Type inference
--> 105     typemap, return_type, calltypes, errs = type_inference_stage(
    106         state.typingctx,
    107         state.targetctx,
    108         state.func_ir,
    109         state.args,
    110         state.return_type,
    111         state.locals,
    112         raise_errors=self._raise_errors)
    113     state.typemap = typemap
    114     # save errors in case of partial typing

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/typed_passes.py:83, in type_inference_stage(typingctx, targetctx, interp, args, return_type, locals, raise_errors)
     81     infer.build_constraint()
     82     # return errors in case of partial typing
---> 83     errs = infer.propagate(raise_errors=raise_errors)
     84     typemap, restype, calltypes = infer.unify(raise_errors=raise_errors)
     86 # Output all Numba warnings

File ~/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/typeinfer.py:1086, in TypeInferer.propagate(self, raise_errors)
   1083 force_lit_args = [e for e in errors
   1084                   if isinstance(e, ForceLiteralArg)]
   1085 if not force_lit_args:
-> 1086     raise errors[0]
   1087 else:
   1088     raise reduce(operator.or_, force_lit_args)

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function numba_funcify_Elemwise.<locals>.elemwise at 0x18dfcdea0>) found for signature:
 
 >>> elemwise(readonly array(float64, 0d, C), readonly array(float64, 0d, C), float64, array(float64, 0d, C))
 
There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload in function 'numba_funcify_Elemwise.<locals>.ov_elemwise': File: pytensor/link/numba/dispatch/elemwise.py: Line 687.
    With argument(s): '(readonly array(float64, 0d, C), readonly array(float64, 0d, C), float64, array(float64, 0d, C))':
   Rejected as the implementation raised a specific error:
     TypingError: Failed in nopython mode pipeline (step: nopython frontend)
   No implementation of function Function(<intrinsic _vectorized>) found for signature:
    
    >>> _vectorized(type(CPUDispatcher(<function numba_funcified_fgraph at 0x18dfcfbe0>)), Literal[str](gASVCAAAAAAAAAAoKSkpKXSULg==
   ), Literal[str](gASVBAAAAAAAAAAphZQu
   ), Literal[str](gASVDQAAAAAAAACMB2Zsb2F0NjSUhZQu
   ), Literal[str](gASVCQAAAAAAAABLAEsChpSFlC4=
   ), StarArgTuple(readonly array(float64, 0d, C), readonly array(float64, 0d, C), float64, array(float64, 0d, C)))
    
   There are 2 candidate implementations:
         - Of which 1 did not match due to:
         Intrinsic in function '_vectorized': File: pytensor/link/numba/dispatch/elemwise.py: Line 466.
           With argument(s): '(type(CPUDispatcher(<function numba_funcified_fgraph at 0x18dfcfbe0>)), unicode_type, unicode_type, unicode_type, unicode_type, StarArgTuple(readonly array(float64, 0d, C), readonly array(float64, 0d, C), float64, array(float64, 0d, C)))':
          Rejected as the implementation raised a specific error:
            TypingError: input_bc_patterns must be literal.
     raised from /Users/bpajusco/opt/anaconda3/envs/fast/lib/python3.10/site-packages/pytensor/link/numba/dispatch/elemwise.py:486
         - Of which 1 did not match due to:
         Intrinsic in function '_vectorized': File: pytensor/link/numba/dispatch/elemwise.py: Line 466.
           With argument(s): '(type(CPUDispatcher(<function numba_funcified_fgraph at 0x18dfcfbe0>)), Literal[str](gASVCAAAAAAAAAAoKSkpKXSULg==
         ), Literal[str](gASVBAAAAAAAAAAphZQu
         ), Literal[str](gASVDQAAAAAAAACMB2Zsb2F0NjSUhZQu
         ), Literal[str](gASVCQAAAAAAAABLAEsChpSFlC4=
         ), StarArgTuple(readonly array(float64, 0d, C), readonly array(float64, 0d, C), float64, array(float64, 0d, C)))':
          Rejected as the implementation raised a specific error:
            TypingError: Inputs to elemwise must be arrays.
     raised from /Users/bpajusco/opt/anaconda3/envs/fast/lib/python3.10/site-packages/pytensor/link/numba/dispatch/elemwise.py:514
   
   During: resolving callee type: Function(<intrinsic _vectorized>)
   During: typing of call at /Users/bpajusco/opt/anaconda3/envs/fast/lib/python3.10/site-packages/pytensor/link/numba/dispatch/elemwise.py (648)
   
   
   File "../../../../opt/anaconda3/envs/fast/lib/python3.10/site-packages/pytensor/link/numba/dispatch/elemwise.py", line 648:
       def elemwise_wrapper(*inputs):
           return _vectorized(
           ^

  raised from /Users/bpajusco/opt/anaconda3/envs/fast/lib/python3.10/site-packages/numba/core/typeinfer.py:1086

During: resolving callee type: Function(<function numba_funcify_Elemwise.<locals>.elemwise at 0x18dfcdea0>)
During: typing of call at /var/folders/z8/q5v42wdj6j1fxt991j966lqm0000gn/T/tmp0kknc9dd (155)


File "../../../../../../var/folders/z8/q5v42wdj6j1fxt991j966lqm0000gn/T/tmp0kknc9dd", line 155:
def numba_funcified_fgraph(_joined_variables):
    <source elided>
    # Elemwise{Composite{((i0 - (i1 * i2)) - i3)}}[(0, 2)](TensorConstant{-234.32932596719152}, TensorConstant{0.5}, InplaceDimShuffle{}.0, Sum{acc_dtype=float64}.0)
    tensor_variable_63 = elemwise_28(tensor_constant_27, tensor_constant_28, tensor_variable_60, tensor_variable_42)
    ^

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x18df54820>))
During: typing of call at /Users/bpajusco/opt/anaconda3/envs/fast/lib/python3.10/site-packages/nutpie/compile_pymc.py (265)

During: resolving callee type: type(CPUDispatcher(<function numba_funcified_fgraph at 0x18df54820>))
During: typing of call at /Users/bpajusco/opt/anaconda3/envs/fast/lib/python3.10/site-packages/nutpie/compile_pymc.py (265)


File "../../../../opt/anaconda3/envs/fast/lib/python3.10/site-packages/nutpie/compile_pymc.py", line 265:
        def extract_shared(x, user_data_):
            return inner(x)
            ^

And also pymc model as a reference:

with pm.Model() as icm:
    # Priors
    λ = pm.Gamma("λ", alpha=2, beta=0.5)
    kernel = pm.gp.cov.ExpQuad(input_dim=2, ls=λ, active_dims=[0])
    σ = pm.HalfNormal("σ", sigma=3)
    
    # Get ICM kernel
    W = pm.Normal("W", mu=0, sigma=3, shape=(n_outputs,2), initval=np.random.randn(n_outputs,2))
    κ = pm.Gamma("κ", alpha=1.5, beta=1, shape=n_outputs)
    B = pm.Deterministic('B', at.dot(W, W.T) + at.diag(κ))
    
    cov_icm = get_icm(input_dim=2, kernel=kernel, B=B, active_dims=[1])

    
    gp = pm.gp.Marginal(cov_func=cov_icm)
    y_ = gp.marginal_likelihood("f", X, Y, sigma=σ)   

[twiecki]

@bpajusco Thanks for reporting. Can you format your message better using ``` ?

[bpajusco]

Done and sorry about not doing it properly the first time around

[aseyboldt]

That looks like an incorrect numba implementation for one of the pytensor ops...
Can you share the get_icm function, so that I can debug locally? If not, it would help to see what the output of this is: pytensor.dprint(icm.logp_dlogp_function()._pytensor_function).

Independent of nutpie vs jax etc, if you want to speed this up, could you maybe use the Woodbury identity to avoid factorizing $κI + WW^T$? (assuming that's what's happening in get_icm)...

[bpajusco]

Here is the output from pytensor.dprint()

Sum{acc_dtype=float64} [id A] '__logp' 115
 |MakeVector{dtype='float64'} [id B] 114
   |Elemwise{Composite{(Switch(i0, ((i1 - (i2 * i3)) + Switch(i4, i5, i6)), i5) + i6)}}[(0, 3)] [id C] 'λ_log___logprob' 113
   | |Elemwise{ge,no_inplace} [id D] 14
   | | |Elemwise{exp,no_inplace} [id E] 'λ_log___log' 1
   | | | |λ_log__ [id F]
   | | |TensorConstant{0.0} [id G]
   | |TensorConstant{-1.3862943611198906} [id H]
   | |TensorConstant{0.5} [id I]
   | |Elemwise{exp,no_inplace} [id E] 'λ_log___log' 1
   | |Elemwise{eq,no_inplace} [id J] 13
   | | |Elemwise{exp,no_inplace} [id E] 'λ_log___log' 1
   | | |TensorConstant{0} [id K]
   | |TensorConstant{-inf} [id L]
   | |λ_log__ [id F]
   |Elemwise{Composite{(Switch(i0, ((i1 - i2) + Switch(i3, i4, (i5 * i6))), i4) + i6)}}[(0, 2)] [id M] 'η_log___logprob' 102
   | |Elemwise{ge,no_inplace} [id N] 11
   | | |Elemwise{exp,no_inplace} [id O] 'η_log___log' 0
   | | | |η_log__ [id P]
   | | |TensorConstant{0.0} [id G]
   | |TensorConstant{-0.6931471805599454} [id Q]
   | |Elemwise{exp,no_inplace} [id O] 'η_log___log' 0
   | |Elemwise{eq,no_inplace} [id R] 10
   | | |Elemwise{exp,no_inplace} [id O] 'η_log___log' 0
   | | |TensorConstant{0} [id K]
   | |TensorConstant{-inf} [id L]
   | |TensorConstant{2.0} [id S]
   | |η_log__ [id P]
   |Elemwise{Composite{(Switch(i0, ((i1 * sqr((i2 * i3))) - i4), i5) + i6)}}[(0, 3)] [id T] 'σ_log___logprob' 101
   | |Elemwise{ge,no_inplace} [id U] 16
   | | |Elemwise{exp,no_inplace} [id V] 'σ_log___log' 2
   | | | |σ_log__ [id W]
   | | |TensorConstant{0.0} [id G]
   | |TensorConstant{-0.5} [id X]
   | |TensorConstant{0.3333333333333333} [id Y]
   | |Elemwise{exp,no_inplace} [id V] 'σ_log___log' 2
   | |TensorConstant{1.3244036413128373} [id Z]
   | |TensorConstant{-inf} [id L]
   | |σ_log__ [id W]
   |Sum{acc_dtype=float64} [id BA] 20
   | |Elemwise{Composite{((i0 * sqr((i1 * i2))) - i3)}} [id BB] 'sigma > 0' 8
   |   |TensorConstant{(1, 1) of -0.5} [id BC]
   |   |TensorConstant{(1, 1) of ..3333333333} [id BD]
   |   |W [id BE]
   |   |TensorConstant{(1, 1) of ..8218727822} [id BF]
   |Sum{acc_dtype=float64} [id BG] 108
   | |Elemwise{Composite{(Switch(i0, ((i1 - i2) + Switch(i3, i4, (i5 * i6))), i4) + i6)}}[(0, 2)] [id BH] 'κ_log___logprob' 106
   |   |Elemwise{ge,no_inplace} [id BI] 19
   |   | |Elemwise{exp,no_inplace} [id BJ] 'κ_log___log' 4
   |   | | |κ_log__ [id BK]
   |   | |TensorConstant{(1,) of 0.0} [id BL]
   |   |TensorConstant{(1,) of 0...3763524526} [id BM]
   |   |Elemwise{exp,no_inplace} [id BJ] 'κ_log___log' 4
   |   |Elemwise{eq,no_inplace} [id BN] 18
   |   | |Elemwise{exp,no_inplace} [id BJ] 'κ_log___log' 4
   |   | |TensorConstant{(1,) of 0} [id BO]
   |   |TensorConstant{(1,) of -inf} [id BP]
   |   |TensorConstant{(1,) of 0.5} [id BQ]
   |   |κ_log__ [id BK]
   |Elemwise{Composite{Switch(i0, ((i1 - (i2 * i3)) - i4), i5)}}[(0, 3)] [id BR] 'posdef' 70
     |All [id BS] 47
     | |Elemwise{gt,no_inplace} [id BT] 43
     |   |ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id BU] 41
     |   | |Cholesky{lower=True, destructive=False, on_error='nan'} [id BV] 39
     |   |   |Elemwise{Composite{((i0 * i1 * i2) + i3 + i4)}} [id BW] 38
     |   |     |Elemwise{Composite{exp((i0 * clip(i1, i2, i3)))}}[(0, 1)] [id BX] 37
     |   |     | |TensorConstant{(1, 1) of -0.5} [id BC]
     |   |     | |Elemwise{Add}[(0, 0)] [id BY] 35
     |   |     | | |Dot22Scalar [id BZ] 31
     |   |     | | | |Elemwise{true_div,no_inplace} [id CA] 22
     |   |     | | | | |TensorConstant{[[ 0.]
 [ ...]
 [84.]]} [id CB]
     |   |     | | | | |InplaceDimShuffle{x,x} [id CC] 12
     |   |     | | | |   |Elemwise{exp,no_inplace} [id E] 'λ_log___log' 1
     |   |     | | | |InplaceDimShuffle{1,0} [id CD] 27
     |   |     | | | | |Elemwise{true_div,no_inplace} [id CA] 22
     |   |     | | | |TensorConstant{-2.0} [id CE]
     |   |     | | |InplaceDimShuffle{0,x} [id CF] 34
     |   |     | | | |Elemwise{sqr,no_inplace} [id CG] 30
     |   |     | | |   |InplaceDimShuffle{0} [id CH] 26
     |   |     | | |     |Elemwise{true_div,no_inplace} [id CA] 22
     |   |     | | |InplaceDimShuffle{x,0} [id CI] 33
     |   |     | |   |Elemwise{sqr,no_inplace} [id CG] 30
     |   |     | |TensorConstant{(1, 1) of 0.0} [id CJ]
     |   |     | |TensorConstant{(1, 1) of inf} [id CK]
     |   |     |Elemwise{sqr,no_inplace} [id CL] 21
     |   |     | |InplaceDimShuffle{x,x} [id CM] 9
     |   |     |   |Elemwise{exp,no_inplace} [id O] 'η_log___log' 0
     |   |     |AdvancedSubtensor [id CN] 32
     |   |     | |SpecifyShape [id CO] 29
     |   |     | | |Gemm{inplace} [id CP] 25
     |   |     | | | |AllocDiag{offset=0, axis1=0, axis2=1} [id CQ] 17
     |   |     | | | | |Elemwise{exp,no_inplace} [id BJ] 'κ_log___log' 4
     |   |     | | | |(d__logp/dη_log___log_jacobian){1.0} [id CR]
     |   |     | | | |W [id BE]
     |   |     | | | |InplaceDimShuffle{1,0} [id CS] 'W.T' 3
     |   |     | | | | |W [id BE]
     |   |     | | | |(d__logp/dη_log___log_jacobian){1.0} [id CR]
     |   |     | | |TensorConstant{3} [id CT]
     |   |     | | |TensorConstant{3} [id CT]
     |   |     | |TensorConstant{[[0]
 [0]
.. [2]
 [2]]} [id CU]
     |   |     | |TensorConstant{[[0 0 0 0 ..
  2 2 2]]} [id CV]
     |   |     |AllocDiag{offset=0, axis1=0, axis2=1} [id CW] 28
     |   |     | |Alloc [id CX] 24
     |   |     |   |Elemwise{sqr,no_inplace} [id CY] 15
     |   |     |   | |Elemwise{exp,no_inplace} [id V] 'σ_log___log' 2
     |   |     |   |TensorConstant{255} [id CZ]
     |   |     |TensorConstant{[[1.e-06 0..0 1.e-06]]} [id DA]
     |   |TensorConstant{(1,) of 0} [id BO]
     |TensorConstant{-234.32932596719152} [id DB]
     |TensorConstant{0.5} [id I]
     |InplaceDimShuffle{} [id DC] 67
     | |Sum{axis=[1], acc_dtype=float64} [id DD] 64
     |   |Elemwise{Sqr}[(0, 0)] [id DE] 62
     |     |InplaceDimShuffle{1,0} [id DF] 58
     |       |SolveTriangular{lower=True, trans=0, unit_diagonal=False, check_finite=True} [id DG] 56
     |         |Elemwise{Switch}[(0, 1)] [id DH] 54
     |         | |InplaceDimShuffle{x,x} [id DI] 51
     |         | | |All [id BS] 47
     |         | |Cholesky{lower=True, destructive=False, on_error='nan'} [id BV] 39
     |         | |TensorConstant{(1, 1) of 1} [id DJ]
     |         |TensorConstant{[[-0.68584..56159017]]} [id DK]
     |Sum{acc_dtype=float64} [id DL] 49
     | |Elemwise{Log}[(0, 0)] [id DM] 45
     |   |ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id BU] 41
     |TensorConstant{-inf} [id L]
Elemwise{Composite{(Switch(i0, (i1 * i2), i3) + Switch(i4, i5, Switch(i0, i6, i5)) + (i7 * i2) + i6)}}[(0, 7)] [id DN] 'λ_log___grad' 112
 |Elemwise{ge,no_inplace} [id D] 14
 |TensorConstant{-0.5} [id X]
 |Elemwise{exp,no_inplace} [id E] 'λ_log___log' 1
 |TensorConstant{0} [id K]
 |Elemwise{eq,no_inplace} [id J] 13
 |TensorConstant{0.0} [id G]
 |(d__logp/dη_log___log_jacobian){1.0} [id CR]
 |InplaceDimShuffle{} [id DO] 111
   |Sum{axis=[0], acc_dtype=float64} [id DP] 110
     |Elemwise{Composite{((-((((i0 * ((i1 * i2 * i3) + (i4 * i2 * i5)) * i6) / i7) + i8) * i6)) / i9)}}[(0, 3)] [id DQ] 109
       |TensorConstant{(1, 1) of 2.0} [id DR]
       |TensorConstant{(1, 1) of -0.5} [id BC]
       |Elemwise{sqr,no_inplace} [id CL] 21
       |InplaceDimShuffle{0,x} [id DS] 99
       | |Sum{axis=[1], acc_dtype=float64} [id DT] 90
       |   |Elemwise{mul,no_inplace} [id DU] 85
       |     |Elemwise{Composite{AND(GE(i0, i1), LE(i0, i2))}} [id DV] 36
       |     | |Elemwise{Add}[(0, 0)] [id BY] 35
       |     | |TensorConstant{(1, 1) of 0.0} [id CJ]
       |     | |TensorConstant{(1, 1) of inf} [id CK]
       |     |Elemwise{Composite{Switch(i0, (((i1 + i2) * i3) - i4), i5)}} [id DW] 81
       |     | |Elemwise{Invert}[(0, 0)] [id DX] 50
       |     | | |InplaceDimShuffle{x,x} [id DY] 46
       |     | |   |Any [id DZ] 42
       |     | |     |Elemwise{isnan,no_inplace} [id EA] 40
       |     | |       |Cholesky{lower=True, destructive=False, on_error='nan'} [id BV] 39
       |     | |SolveTriangular{lower=False, trans=0, unit_diagonal=False, check_finite=True} [id EB] 77
       |     | | |InplaceDimShuffle{1,0} [id EC] 55
       |     | | | |Elemwise{switch,no_inplace} [id ED] 53
       |     | | |   |Elemwise{Invert}[(0, 0)] [id DX] 50
       |     | | |   |Cholesky{lower=True, destructive=False, on_error='nan'} [id BV] 39
       |     | | |   |TensorConstant{(1, 1) of 1} [id DJ]
       |     | | |InplaceDimShuffle{1,0} [id EE] 76
       |     | |   |SolveTriangular{lower=False, trans=0, unit_diagonal=False, check_finite=True} [id EF] 75
       |     | |     |InplaceDimShuffle{1,0} [id EC] 55
       |     | |     |Elemwise{Composite{((i0 * i1) - i2)}}[(0, 0)] [id EG] 74
       |     | |       |Dot22 [id EH] 69
       |     | |       | |InplaceDimShuffle{1,0} [id EI] 66
       |     | |       | | |Elemwise{Composite{Switch(i0, (i1 + Switch(i2, (i3 * i4), i5)), i6)}}[(0, 1)] [id EJ] 63
       |     | |       | |   |Elemwise{Invert}[(0, 0)] [id DX] 50
       |     | |       | |   |IncSubtensor{InplaceSet;:int64:, :int64:} [id EK] 52
       |     | |       | |   | |Alloc [id EL] 5
       |     | |       | |   | | |TensorConstant{(1, 1) of 0.0} [id EM]
       |     | |       | |   | | |TensorConstant{255} [id EN]
       |     | |       | |   | | |TensorConstant{255} [id EN]
       |     | |       | |   | |AllocDiag{offset=0, axis1=0, axis2=1} [id EO] 48
       |     | |       | |   | | |Elemwise{true_div} [id EP] 44
       |     | |       | |   | |   |TensorConstant{(1,) of -1.0} [id EQ]
       |     | |       | |   | |   |ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id BU] 41
       |     | |       | |   | |ScalarConstant{255} [id ER]
       |     | |       | |   | |ScalarConstant{255} [id ER]
       |     | |       | |   |InplaceDimShuffle{x,x} [id DI] 51
       |     | |       | |   |Dot22Scalar [id ES] 61
       |     | |       | |   | |SolveTriangular{lower=False, trans=0, unit_diagonal=False, check_finite=True} [id ET] 60
       |     | |       | |   | | |InplaceDimShuffle{1,0} [id EU] 57
       |     | |       | |   | | | |Elemwise{Switch}[(0, 1)] [id DH] 54
       |     | |       | |   | | |Elemwise{mul,no_inplace} [id EV] 59
       |     | |       | |   | |   |TensorConstant{(1, 1) of 2.0} [id DR]
       |     | |       | |   | |   |TensorConstant{(1, 1) of -0.5} [id BC]
       |     | |       | |   | |   |SolveTriangular{lower=True, trans=0, unit_diagonal=False, check_finite=True} [id DG] 56
       |     | |       | |   | |InplaceDimShuffle{1,0} [id DF] 58
       |     | |       | |   | |TensorConstant{-1.0} [id EW]
       |     | |       | |   |TensorConstant{[[1. 0. 0...1. 1. 1.]]} [id EX]
       |     | |       | |   |TensorConstant{(1, 1) of 0.0} [id CJ]
       |     | |       | |   |TensorConstant{(1, 1) of 1} [id DJ]
       |     | |       | |Elemwise{switch,no_inplace} [id ED] 53
       |     | |       |TensorConstant{[[1. 1. 1...0. 0. 1.]]} [id EY]
       |     | |       |InplaceDimShuffle{1,0} [id EZ] 73
       |     | |         |AllocDiag{offset=0, axis1=0, axis2=1} [id FA] 72
       |     | |           |Elemwise{Mul}[(0, 1)] [id FB] 71
       |     | |             |TensorConstant{(1,) of 0.5} [id BQ]
       |     | |             |ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id FC] 68
       |     | |               |Dot22 [id FD] 65
       |     | |                 |InplaceDimShuffle{1,0} [id EC] 55
       |     | |                 |Elemwise{Composite{Switch(i0, (i1 + Switch(i2, (i3 * i4), i5)), i6)}}[(0, 1)] [id EJ] 63
       |     | |InplaceDimShuffle{1,0} [id FE] 79
       |     | | |SolveTriangular{lower=False, trans=0, unit_diagonal=False, check_finite=True} [id EB] 77
       |     | |TensorConstant{[[1. 0. 0...1. 1. 1.]]} [id EX]
       |     | |AllocDiag{offset=0, axis1=0, axis2=1} [id FF] 80
       |     | | |ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id FG] 78
       |     | |   |SolveTriangular{lower=False, trans=0, unit_diagonal=False, check_finite=True} [id EB] 77
       |     | |TensorConstant{[[nan]]} [id FH]
       |     |AdvancedSubtensor [id CN] 32
       |     |Elemwise{Composite{exp((i0 * clip(i1, i2, i3)))}}[(0, 1)] [id BX] 37
       |TensorConstant{(1, 1) of -0.5} [id BC]
       |InplaceDimShuffle{0,x} [id FI] 100
       | |Sum{axis=[0], acc_dtype=float64} [id FJ] 91
       |   |Elemwise{mul,no_inplace} [id DU] 85
       |TensorConstant{[[ 0.]
 [ ...]
 [84.]]} [id CB]
       |InplaceDimShuffle{x,x} [id CC] 12
       |InplaceDimShuffle{0,x} [id FK] 107
       | |CGemv{inplace} [id FL] 103
       |   |CGemv{inplace} [id FM] 93
       |   | |AllocEmpty{dtype='float64'} [id FN] 7
       |   | | |TensorConstant{255} [id FO]
       |   | |(d__logp/dη_log___log_jacobian){1.0} [id CR]
       |   | |Elemwise{Mul}[(0, 1)] [id FP] 86
       |   | | |Elemwise{Composite{AND(GE(i0, i1), LE(i0, i2))}} [id DV] 36
       |   | | |Elemwise{Composite{Switch(i0, (((i1 + i2) * i3) - i4), i5)}} [id DW] 81
       |   | | |Elemwise{sqr,no_inplace} [id CL] 21
       |   | | |AdvancedSubtensor [id CN] 32
       |   | | |Elemwise{Composite{exp((i0 * clip(i1, i2, i3)))}}[(0, 1)] [id BX] 37
       |   | |InplaceDimShuffle{0} [id CH] 26
       |   | |TensorConstant{0.0} [id FQ]
       |   |(d__logp/dη_log___log_jacobian){1.0} [id CR]
       |   |InplaceDimShuffle{1,0} [id FR] 92
       |   | |Elemwise{Mul}[(0, 1)] [id FP] 86
       |   |InplaceDimShuffle{0} [id CH] 26
       |   |(d__logp/dη_log___log_jacobian){1.0} [id CR]
       |Elemwise{sqr,no_inplace} [id FS] 23
         |InplaceDimShuffle{x,x} [id CC] 12
Elemwise{Composite{(Switch(i0, (-i1), i2) + Switch(i3, i2, Switch(i0, i4, i2)) + (i5 * i6 * i1 * i1) + i7)}}[(0, 6)] [id FT] 'η_log___grad' 98
 |Elemwise{ge,no_inplace} [id N] 11
 |Elemwise{exp,no_inplace} [id O] 'η_log___log' 0
 |TensorConstant{0} [id K]
 |Elemwise{eq,no_inplace} [id R] 10
 |TensorConstant{2.0} [id S]
 |TensorConstant{2.0} [id S]
 |Sum{acc_dtype=float64} [id FU] 89
 | |Elemwise{mul} [id FV] 84
 |   |Elemwise{Composite{Switch(i0, (((i1 + i2) * i3) - i4), i5)}} [id DW] 81
 |   |Elemwise{Composite{exp((i0 * clip(i1, i2, i3)))}}[(0, 1)] [id BX] 37
 |   |AdvancedSubtensor [id CN] 32
 |(d__logp/dη_log___log_jacobian){1.0} [id CR]
Elemwise{Composite{(Switch(i0, (i1 * i2 * i2), i3) + (i4 * i5 * i2 * i2) + i6)}}[(0, 5)] [id FW] 'σ_log___grad' 97
 |Elemwise{ge,no_inplace} [id U] 16
 |TensorConstant{-0.1111111111111111} [id FX]
 |Elemwise{exp,no_inplace} [id V] 'σ_log___log' 2
 |TensorConstant{0} [id K]
 |TensorConstant{2.0} [id S]
 |Sum{acc_dtype=float64} [id FY] 88
 | |ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id FZ] 83
 |   |Elemwise{Composite{Switch(i0, (((i1 + i2) * i3) - i4), i5)}} [id DW] 81
 |(d__logp/dη_log___log_jacobian){1.0} [id CR]
Gemm{inplace} [id GA] 'W_grad' 105
 |Gemm{no_inplace} [id GB] 96
 | |W [id BE]
 | |(d__logp/dη_log___log_jacobian){1.0} [id CR]
 | |AdvancedIncSubtensor{inplace=True,  set_instead_of_inc=False} [id GC] 87
 | | |Alloc [id GD] 6
 | | | |TensorConstant{(1, 1) of 0.0} [id EM]
 | | | |TensorConstant{3} [id GE]
 | | | |TensorConstant{3} [id GE]
 | | |Elemwise{mul} [id GF] 82
 | | | |Elemwise{Composite{Switch(i0, (((i1 + i2) * i3) - i4), i5)}} [id DW] 81
 | | | |Elemwise{Composite{exp((i0 * clip(i1, i2, i3)))}}[(0, 1)] [id BX] 37
 | | | |Elemwise{sqr,no_inplace} [id CL] 21
 | | |TensorConstant{[[0]
 [0]
.. [2]
 [2]]} [id CU]
 | | |TensorConstant{[[0 0 0 0 ..
  2 2 2]]} [id CV]
 | |W [id BE]
 | |TensorConstant{-0.1111111111111111} [id FX]
 |(d__logp/dη_log___log_jacobian){1.0} [id CR]
 |InplaceDimShuffle{1,0} [id GG] 95
 | |AdvancedIncSubtensor{inplace=True,  set_instead_of_inc=False} [id GC] 87
 |W [id BE]
 |(d__logp/dη_log___log_jacobian){1.0} [id CR]
Elemwise{Composite{(Switch(i0, (-i1), i2) + Switch(i3, i2, Switch(i0, i4, i2)) + (i5 * i1) + i6)}} [id GH] 'κ_log___grad' 104
 |Elemwise{ge,no_inplace} [id BI] 19
 |Elemwise{exp,no_inplace} [id BJ] 'κ_log___log' 4
 |TensorConstant{(1,) of 0} [id BO]
 |Elemwise{eq,no_inplace} [id BN] 18
 |TensorConstant{(1,) of 0.5} [id BQ]
 |ExtractDiag{offset=0, axis1=0, axis2=1, view=False} [id GI] 94
 | |AdvancedIncSubtensor{inplace=True,  set_instead_of_inc=False} [id GC] 87
 |TensorConstant{(1,) of 1.0} [id GJ]

and here is get_icm function:

def get_icm(input_dim, kernel, W=None, kappa=None, B=None, active_dims=None, name='ICM'):
    coreg = pm.gp.cov.Coregion(input_dim=input_dim, W=W, kappa=kappa, B=B, active_dims=active_dims)
    icm_cov = kernel * coreg 
    return icm_cov

And indeed this is what get_icm does

[aseyboldt]

@bpajusco Should be fixed on pytensor main now.
You still won't get great performance however, because there are still two Ops in your graph, (SolveTriangular and AdvancedSubtensor) that do not have a numba implementation and need to call into python.

@bwengals I was looking into the Coregion covariance (essentially $I + WW^T$, where $W$ is low-rank), and it doesn't seem we are taking advantage of the low rank structure. Do you think it would make sense to extend BaseCovariance, adding methods like those, that have a default implementation using factorizations, but could be overwritten by special covariance where we have better ways to do the computation?

Something along the lines of (with better names?)

• logdet_and_solve: Compute logdet(C) and C^{-1}x. This could then be used to implement the logp.
• factorization_product: Compute $Zx$ where $C = ZZ^T$ is a factorization. This could be used for non-centered parametrizations.
• matvec: $Cx$. Could for instance be used to iteratively compute $sqrt(C)x$?

[bwengals]

Hey thanks for tagging me, super interested in this. I resp here: pymc-devs/pymc#6615