How to use the exoplanet.theano_ops.celerite.base_op.CeleriteBaseOp function in exoplanet

# -*- coding: utf-8 -*-

__all__ = ["SolveRevOp"]

import theano.tensor as tt
from theano import gof

from .base_op import CeleriteBaseOp


class SolveRevOp(CeleriteBaseOp):

    func_file = "./solve_rev.cc"
    func_name = "APPLY_SPECIFIC(solve_rev)"

    def make_node(self, *args):
        in_args = [tt.as_tensor_variable(a) for a in args]
        out_args = [a.type() for a in args[:5]]
        return gof.Apply(self, in_args, out_args)

# -*- coding: utf-8 -*-

__all__ = ["FactorRevOp"]

import theano.tensor as tt
from theano import gof

from .base_op import CeleriteBaseOp


class FactorRevOp(CeleriteBaseOp):

    func_file = "./factor_rev.cc"
    func_name = "APPLY_SPECIFIC(factor_rev)"

    def __init__(self, J=-1):
        super(FactorRevOp, self).__init__(J=J)

    def make_node(self, *args):
        in_args = [tt.as_tensor_variable(a) for a in args]
        out_args = [
            in_args[2].type(),
            in_args[0].type(),
            in_args[3].type(),
            in_args[1].type(),
        ]
        return gof.Apply(self, in_args, out_args)

# -*- coding: utf-8 -*-

__all__ = ["DiagDotOp"]

import theano
import theano.tensor as tt
from theano import gof

from .base_op import CeleriteBaseOp


class DiagDotOp(CeleriteBaseOp):

    func_file = "./diag_dot.cc"
    func_name = "APPLY_SPECIFIC(diag_dot)"

    def make_node(self, *args):
        in_args = [tt.as_tensor_variable(a) for a in args]
        out_args = [tt.vector(dtype=theano.config.floatX).type()]
        return gof.Apply(self, in_args, out_args)

    def grad(self, inputs, gradients):
        A, B = inputs
        return (
            gradients[0][:, None] * tt.transpose(B),
            gradients[1][None, :] * tt.transpose(A),
        )

# -*- coding: utf-8 -*-

__all__ = ["SolveOp"]

import theano
import theano.tensor as tt
from theano import gof

from .base_op import CeleriteBaseOp
from .solve_rev import SolveRevOp


class SolveOp(CeleriteBaseOp):

    func_file = "./solve.cc"
    func_name = "APPLY_SPECIFIC(solve)"
    num_input = 5
    output_ndim = (2, 2, 2)

    def __init__(self, J=-1, n_rhs=-1):
        self.grad_op = SolveRevOp(J=J, n_rhs=n_rhs)
        super(SolveOp, self).__init__(J=J, n_rhs=n_rhs)

    def make_node(self, *args):
        in_args = [tt.as_tensor_variable(a) for a in args]
        out_args = [
            in_args[-1].type(),
            tt.matrix(dtype=theano.config.floatX).type(),
            tt.matrix(dtype=theano.config.floatX).type(),

def __init__(self, J=-1, n_rhs=-1):
        self.J = int(J)
        self.n_rhs = int(n_rhs)
        super(CeleriteBaseOp, self).__init__(self.func_file, self.func_name)

# -*- coding: utf-8 -*-

__all__ = ["FactorOp"]

import theano
import theano.tensor as tt
from theano import gof

from .base_op import CeleriteBaseOp
from .factor_rev import FactorRevOp


class FactorOp(CeleriteBaseOp):

    func_file = "./factor.cc"
    func_name = "APPLY_SPECIFIC(factor)"

    def __init__(self, J=-1):
        self.grad_op = FactorRevOp(J=J)
        super(FactorOp, self).__init__(J=J)

    def make_node(self, *args):
        in_args = [tt.as_tensor_variable(a) for a in args]
        out_args = [
            in_args[0].type(),
            in_args[2].type(),
            tt.matrix(dtype=theano.config.floatX).type(),
            tt.iscalar().type(),
        ]