"""
ML Autogen
==========

This is an autogenerated file DO NOT TOUCH, it's sole intention is to provide IDEs
a way to index functions and thereby supporting the users with autocompletion.

Generated by: autogen.py
Complements: on_ml.py

Time of generation (UTC): 2022-03-21T07:12:52.019635
"""
# -----------------------------------------------------------------------------------

# The classes here are designed like gRPC stubs, templates that can be indexed by the
# IDE and intuitively used by the users by just filling the variable in it.

from typing import Union, Any
from google.protobuf.timestamp_pb2 import Timestamp

from .model_spec_pb2 import ModelSpec, Framework
from .on_ml import NboxOptions

# --------- torch --------- #

class TorchToTorchscript:
  # @isthere("-f https://download.pytorch.org/whl/cpu/torch_stable.html","torch", soft=False)
  def __init__(
    self,
    func=None,
    example_inputs=None,
    optimize=None,
    check_trace=True,
    check_inputs=None,
    check_tolerance=1e-05,
    strict=True,
    _force_outplace=False,
    _module_class=None,
    **kwargs
  ):
    '''AUTOGENERATED by ``nbox``, some arguments might be missing, USE ``kwargs`` CAREFULLY!

    Trace a function and return an executable  or :class:`ScriptFunction`
    that will be optimized using just-in-time compilation. Tracing is ideal for
    code that operates only on ``Tensor``\s and lists, dictionaries, and
    tuples of ``Tensor``\s.

    Using `torch.jit.trace` and `torch.jit.trace_module`, you can turn an
    existing module or Python function into a TorchScript
    :class:`ScriptFunction` or :class:`ScriptModule`. You must provide example
    inputs, and we run the function, recording the operations performed on all
    the tensors.

    * The resulting recording of a standalone function produces `ScriptFunction`.
    * The resulting recording of `nn.Module.forward` or `nn.Module` produces
      `ScriptModule`.

    This module also contains any parameters that the original
    module had as well.

    Warning:
        Tracing only correctly records functions and modules which are not data
        dependent (e.g., do not have conditionals on data in tensors) and do not have
        any untracked external dependencies (e.g., perform input/output or
        access global variables). Tracing only records operations done when the given
        function is run on the given tensors. Therefore, the returned
        `ScriptModule` will always run the same traced graph on any input. This
        has some important implications when your module is expected to run
        different sets of operations, depending on the input and/or the module
        state. For example,

        * Tracing will not record any control-flow like if-statements or loops.
          When this control-flow is constant across your module, this is fine
          and it often inlines the control-flow decisions. But sometimes the
          control-flow is actually part of the model itself. For instance, a
          recurrent network is a loop over the (possibly dynamic) length of an
          input sequence.
        * In the returned :class:`ScriptModule`, operations that have different
          behaviors in ``training`` and ``eval`` modes will always behave as if
          it is in the mode it was in during tracing, no matter which mode the
          `ScriptModule` is in.

        In cases like these, tracing would not be appropriate and
        :func:`scripting <torch.jit.script>` is a better choice. If you trace
        such models, you may silently get incorrect results on subsequent
        invocations of the model. The tracer will try to emit warnings when
        doing something that may cause an incorrect trace to be produced.

    Args:
        func (callable or torch.nn.Module):  A Python function or `torch.nn.Module`
            that will be run with `example_inputs`. `func` arguments and return
            values  must be tensors or (possibly nested) tuples that contain
            tensors. When a module is passed `torch.jit.trace`, only the
            ``forward`` method is run and traced (see :func:`torch.jit.trace
            <torch.jit.trace_module>` for details).
        example_inputs (tuple or torch.Tensor):  A tuple of example inputs that
            will be passed to the function while tracing. The resulting trace
            can be run with inputs of different types and shapes assuming the
            traced operations support those types and shapes. `example_inputs`
            may also be a single Tensor in which case it is automatically
            wrapped in a tuple.

    Keyword arguments:
        check_trace (``bool``, optional): Check if the same inputs run through
            traced code produce the same outputs. Default: ``True``. You might want
            to disable this if, for example, your network contains non-
            deterministic ops or if you are sure that the network is correct despite
            a checker failure.

        check_inputs (list of tuples, optional): A list of tuples of input
            arguments that should be used to check the trace against what is
            expected. Each tuple is equivalent to a set of input arguments that
            would be specified in ``example_inputs``. For best results, pass in
            a set of checking inputs representative of the space of shapes and
            types of inputs you expect the network to see.  If not specified,
            the original ``example_inputs`` are used for checking
        check_tolerance (float, optional): Floating-point comparison tolerance
            to use in the checker procedure.  This can be used to relax the
            checker strictness in the event that results diverge numerically
            for a known reason, such as operator fusion.
        strict (``bool``, optional): run the tracer in a strict mode or not
            (default: ``True``). Only turn this off when you want the tracer to
            record your mutable container types (currently ``list``/``dict``)
            and you are sure that the container you are using in your
            problem is a ``constant`` structure and does not get used as
            control flow (if, for) conditions.

    Returns:
        If `func` is `nn.Module` or ``forward`` of `nn.Module`, `trace` returns
        a :class:`ScriptModule` object with a single ``forward`` method
        containing the traced code.  The returned `ScriptModule` will
        have the same set of sub-modules and parameters as the original
        ``nn.Module``.  If ``func`` is a standalone function, ``trace``
        returns `ScriptFunction`.

    Example (tracing a function):

    .. testcode::

        import torch

        def foo(x, y):
            return 2 * x + y

        # Run `foo` with the provided inputs and record the tensor operations
        traced_foo = torch.jit.trace(foo, (torch.rand(3), torch.rand(3)))

        # `traced_foo` can now be run with the TorchScript interpreter or saved
        # and loaded in a Python-free environment

    Example (tracing an existing module)::

        import torch
        import torch.nn as nn

        class Net(nn.Module):
            def __init__(self):
                super(Net, self).__init__()
                self.conv = nn.Conv2d(1, 1, 3)

            def forward(self, x):
                return self.conv(x)

        n = Net()
        example_weight = torch.rand(1, 1, 3, 3)
        example_forward_input = torch.rand(1, 1, 3, 3)

        # Trace a specific method and construct `ScriptModule` with
        # a single `forward` method
        module = torch.jit.trace(n.forward, example_forward_input)

        # Trace a module (implicitly traces `forward`) and construct a
        # `ScriptModule` with a single `forward` method
        module = torch.jit.trace(n, example_forward_input)
    '''
    self.func=func
    self.example_inputs=example_inputs
    self.optimize=optimize
    self.check_trace=check_trace
    self.check_inputs=check_inputs
    self.check_tolerance=check_tolerance
    self.strict=strict
    self._force_outplace=_force_outplace
    self._module_class=_module_class
    
    for k,v in kwargs.items():
      setattr(self, k, v)

class TorchscriptToTorch:
  # @isthere("-f https://download.pytorch.org/whl/cpu/torch_stable.html","torch", soft=False)
  def __init__(
    self,
    map_location=None,
    _extra_files=None,
    **kwargs
  ):
    '''AUTOGENERATED by ``nbox``, some arguments might be missing, USE ``kwargs`` CAREFULLY!

    Load a :class:`ScriptModule` or :class:`ScriptFunction` previously
    saved with :func:`torch.jit.save <torch.jit.save>`

    All previously saved modules, no matter their device, are first loaded onto CPU,
    and then are moved to the devices they were saved from. If this fails (e.g.
    because the run time system doesn't have certain devices), an exception is
    raised.

    Args:
        f: a file-like object (has to implement read, readline, tell, and seek),
            or a string containing a file name
        map_location (string or torch.device): A simplified version of
            ``map_location`` in `torch.jit.save` used to dynamically remap
            storages to an alternative set of devices.
        _extra_files (dictionary of filename to content): The extra
            filenames given in the map would be loaded and their content
            would be stored in the provided map.

    Returns:
        A :class:`ScriptModule` object.

    Example:

    .. testcode::

        import torch
        import io

        torch.jit.load('scriptmodule.pt')

        # Load ScriptModule from io.BytesIO object
        with open('scriptmodule.pt', 'rb') as f:
            buffer = io.BytesIO(f.read())

        # Load all tensors to the original device
        torch.jit.load(buffer)

        # Load all tensors onto CPU, using a device
        buffer.seek(0)
        torch.jit.load(buffer, map_location=torch.device('cpu'))

        # Load all tensors onto CPU, using a string
        buffer.seek(0)
        torch.jit.load(buffer, map_location='cpu')

        # Load with extra files.
        extra_files = {'foo.txt': ''}  # values will be replaced with data
        torch.jit.load('scriptmodule.pt', _extra_files=extra_files)
        print(extra_files['foo.txt'])

    .. testoutput::
        :hide:

        ...

    .. testcleanup::

        import os
        os.remove("scriptmodule.pt")
    '''
    self.map_location=map_location
    self._extra_files=_extra_files
    
    for k,v in kwargs.items():
      setattr(self, k, v)

class TorchToOnnx:
  # @isthere("-f https://download.pytorch.org/whl/cpu/torch_stable.html","torch", soft=False)
  def __init__(
    self,
    model=None,
    args=None,
    export_params=True,
    verbose=False,
    input_names=None,
    output_names=None,
    operator_export_type=None,
    opset_version=None,
    do_constant_folding=True,
    dynamic_axes=None,
    keep_initializers_as_inputs=None,
    custom_opsets=None,
    export_modules_as_functions=False,
    **kwargs
  ):
    '''AUTOGENERATED by ``nbox``, some arguments might be missing, USE ``kwargs`` CAREFULLY!

    Exports a model into ONNX format. If ``model`` is not a
    :class:`torch.jit.ScriptModule` nor a :class:`torch.jit.ScriptFunction`, this runs
    ``model`` once in order to convert it to a TorchScript graph to be exported
    (the equivalent of :func:`torch.jit.trace`). Thus this has the same limited support
    for dynamic control flow as :func:`torch.jit.trace`.

    Args:
        model (torch.nn.Module, torch.jit.ScriptModule or torch.jit.ScriptFunction):
            the model to be exported.
        args (tuple or torch.Tensor):

            args can be structured either as:

            1. ONLY A TUPLE OF ARGUMENTS::

                args = (x, y, z)

            The tuple should contain model inputs such that ``model(*args)`` is a valid
            invocation of the model. Any non-Tensor arguments will be hard-coded into the
            exported model; any Tensor arguments will become inputs of the exported model,
            in the order they occur in the tuple.

            2. A TENSOR::

                args = torch.Tensor([1])

            This is equivalent to a 1-ary tuple of that Tensor.

            3. A TUPLE OF ARGUMENTS ENDING WITH A DICTIONARY OF NAMED ARGUMENTS::

                args = (x,
                        {'y': input_y,
                         'z': input_z})

            All but the last element of the tuple will be passed as non-keyword arguments,
            and named arguments will be set from the last element. If a named argument is
            not present in the dictionary, it is assigned the default value, or None if a
            default value is not provided.

            .. note::
                If a dictionary is the last element of the args tuple, it will be
                interpreted as containing named arguments. In order to pass a dict as the
                last non-keyword arg, provide an empty dict as the last element of the args
                tuple. For example, instead of::

                    torch.onnx.export(
                        model,
                        (x,
                         # WRONG: will be interpreted as named arguments
                         {y: z}),
                        "test.onnx.pb")

                Write::

                    torch.onnx.export(
                        model,
                        (x,
                         {y: z},
                         {}),
                        "test.onnx.pb")

        f: a file-like object (such that ``f.fileno()`` returns a file descriptor)
            or a string containing a file name.  A binary protocol buffer will be written
            to this file.
        export_params (bool, default True): if True, all parameters will
            be exported. Set this to False if you want to export an untrained model.
            In this case, the exported model will first take all of its parameters
            as arguments, with the ordering as specified by ``model.state_dict().values()``
        verbose (bool, default False): if True, prints a description of the
            model being exported to stdout. In addition, the final ONNX graph will include the
            field ``doc_string``` from the exported model which mentions the source code locations
            for ``model``.
        training (enum, default TrainingMode.EVAL):
            * ``TrainingMode.EVAL``: export the model in inference mode.
            * ``TrainingMode.PRESERVE``: export the model in inference mode if model.training is
              False and in training mode if model.training is True.
            * ``TrainingMode.TRAINING``: export the model in training mode. Disables optimizations
              which might interfere with training.
        input_names (list of str, default empty list): names to assign to the
            input nodes of the graph, in order.
        output_names (list of str, default empty list): names to assign to the
            output nodes of the graph, in order.
        operator_export_type (enum, default None):

            None usually means ``OperatorExportTypes.ONNX``.
            However if PyTorch was built with ``-DPYTORCH_ONNX_CAFFE2_BUNDLE``, None means
            ``OperatorExportTypes.ONNX_ATEN_FALLBACK``.

            * ``OperatorExportTypes.ONNX``: Export all ops as regular ONNX ops
              (in the default opset domain).
            * ``OperatorExportTypes.ONNX_FALLTHROUGH``: Try to convert all ops
              to standard ONNX ops in the default opset domain. If unable to do so
              (e.g. because support has not been added to convert a particular torch op to ONNX),
              fall back to exporting the op into a custom opset domain without conversion. Applies
              to `custom ops <https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html>`_
              as well as ATen ops. For the exported model to be usable, the runtime must support
              these non-standard ops.
            * ``OperatorExportTypes.ONNX_ATEN``: All ATen ops (in the TorchScript namespace "aten")
              are exported as ATen ops (in opset domain "org.pytorch.aten").
              `ATen <https://pytorch.org/cppdocs/#aten>`_ is PyTorch's built-in tensor library, so
              this instructs the runtime to use PyTorch's implementation of these ops.

              .. warning::

                Models exported this way are probably runnable only by Caffe2.

              This may be useful if the numeric differences in implementations of operators are
              causing large differences in behavior between PyTorch and Caffe2 (which is more
              common on untrained models).

            * ``OperatorExportTypes.ONNX_ATEN_FALLBACK``: Try to export each ATen op
              (in the TorchScript namespace "aten") as a regular ONNX op. If we are unable to do so
              (e.g. because support has not been added to convert a particular torch op to ONNX),
              fall back to exporting an ATen op. See documentation on OperatorExportTypes.ONNX_ATEN for
              context.
              For example::

                graph(%0 : Float):
                  %3 : int = prim::Constant[value=0]()
                  # conversion unsupported
                  %4 : Float = aten::triu(%0, %3)
                  # conversion supported
                  %5 : Float = aten::mul(%4, %0)
                  return (%5)

              Assuming ``aten::triu`` is not supported in ONNX, this will be exported as::

                graph(%0 : Float):
                  %1 : Long() = onnx::Constant[value={0}]()
                  # not converted
                  %2 : Float = aten::ATen[operator="triu"](%0, %1)
                  # converted
                  %3 : Float = onnx::Mul(%2, %0)
                  return (%3)

              If PyTorch was built with Caffe2 (i.e. with ``BUILD_CAFFE2=1``), then
              Caffe2-specific behavior will be enabled, including special support
              for ops are produced by the modules described in
              `Quantization <https://pytorch.org/docs/stable/quantization.html>`_.

              .. warning::

                Models exported this way are probably runnable only by Caffe2.

        opset_version (int, default 9): The version of the
            `default (ai.onnx) opset <https://github.com/onnx/onnx/blob/master/docs/Operators.md>`_
            to target. Must be >= 7 and <= 15.
        do_constant_folding (bool, default True): Apply the constant-folding optimization.
            Constant-folding will replace some of the ops that have all constant inputs
            with pre-computed constant nodes.
        dynamic_axes (dict<string, dict<int, string>> or dict<string, list(int)>, default empty dict):

            By default the exported model will have the shapes of all input and output tensors
            set to exactly match those given in ``args``. To specify axes of tensors as
            dynamic (i.e. known only at run-time), set ``dynamic_axes`` to a dict with schema:

            * KEY (str): an input or output name. Each name must also be provided in ``input_names`` or
              ``output_names``.
            * VALUE (dict or list): If a dict, keys are axis indices and values are axis names. If a
              list, each element is an axis index.

            For example::

                class SumModule(torch.nn.Module):
                    def forward(self, x):
                        return torch.sum(x, dim=1)

                torch.onnx.export(SumModule(), (torch.ones(2, 2),), "onnx.pb",
                                  input_names=["x"], output_names=["sum"])

            Produces::

                input {
                  name: "x"
                  ...
                      shape {
                        dim {
                          dim_value: 2  # axis 0
                        }
                        dim {
                          dim_value: 2  # axis 1
                ...
                output {
                  name: "sum"
                  ...
                      shape {
                        dim {
                          dim_value: 2  # axis 0
                ...

            While::

                torch.onnx.export(SumModule(), (torch.ones(2, 2),), "onnx.pb",
                                  input_names=["x"], output_names=["sum"],
                                  dynamic_axes={
                                      # dict value: manually named axes
                                      "x": {0: "my_custom_axis_name"},
                                      # list value: automatic names
                                      "sum": [0],
                                  })

            Produces::

                input {
                  name: "x"
                  ...
                      shape {
                        dim {
                          dim_param: "my_custom_axis_name"  # axis 0
                        }
                        dim {
                          dim_value: 2  # axis 1
                ...
                output {
                  name: "sum"
                  ...
                      shape {
                        dim {
                          dim_param: "sum_dynamic_axes_1"  # axis 0
                ...

        keep_initializers_as_inputs (bool, default None): If True, all the
            initializers (typically corresponding to parameters) in the
            exported graph will also be added as inputs to the graph. If False,
            then initializers are not added as inputs to the graph, and only
            the non-parameter inputs are added as inputs.
            This may allow for better optimizations (e.g. constant folding) by
            backends/runtimes.

            If ``opset_version < 9``, initializers MUST be part of graph
            inputs and this argument will be ignored and the behavior will be
            equivalent to setting this argument to True.

            If None, then the behavior is chosen automatically as follows:

            * If ``operator_export_type=OperatorExportTypes.ONNX``, the behavior is equivalent
              to setting this argument to False.
            * Else, the behavior is equivalent to setting this argument to True.

        custom_opsets (dict<str, int>, default empty dict): A dict with schema:

            * KEY (str): opset domain name
            * VALUE (int): opset version

            If a custom opset is referenced by ``model`` but not mentioned in this dictionary,
            the opset version is set to 1. Only custom opset domain name and version should be
            indicated through this argument.

        export_modules_as_functions (bool or set of type of nn.Module, default False): Flag to enable
            exporting all ``nn.Module`` forward calls as local functions in ONNX. Or a set to indicate the
            particular types of modules to export as local functions in ONNX.
            This feature requires ``opset_version`` >= 15, otherwise the export will fail. This is because
            ``opset_version`` < 15 implies IR version < 8, which means no local function support.

            * ``False``(default): export ``nn.Module`` forward calls as fine grained nodes.
            * ``True``: export all ``nn.Module`` forward calls as local function nodes.
            * Set of type of nn.Module: export ``nn.Module`` forward calls as local function nodes,
              only if the type of the ``nn.Module`` is found in the set.

    Raises:
      CheckerError: If the ONNX checker detects an invalid ONNX graph. Will still export the
        model to the file ``f`` even if this is raised.
    '''
    self.model=model
    self.args=args
    self.export_params=export_params
    self.verbose=verbose
    self.input_names=input_names
    self.output_names=output_names
    self.operator_export_type=operator_export_type
    self.opset_version=opset_version
    self.do_constant_folding=do_constant_folding
    self.dynamic_axes=dynamic_axes
    self.keep_initializers_as_inputs=keep_initializers_as_inputs
    self.custom_opsets=custom_opsets
    self.export_modules_as_functions=export_modules_as_functions
    
    for k,v in kwargs.items():
      setattr(self, k, v)


class Framework_torch:
  _load_framework="torch"
  def _conditional(py_model):
    from .on_ml import torch_condition
    return torch_condition(py_model)
  
  # @isthere("-f https://download.pytorch.org/whl/cpu/torch_stable.html","torch", soft=False)
  def torch_to_torchscript(user_options: TorchToTorchscript, nbox_options: NboxOptions) -> Union[ModelSpec, Any]:
    '''
    Args:
      user_options (nbox.framework.ml.TorchToTorchscript): Variables to be filled by the user
      _nbox_options (nbox.framework.NboxOptions): Options to be used by NBX-Services
    Returns:
      when serialising the model returns ``ModelSpec`` else returns the ``model`` in ``nbox.Model``
    '''
    from .on_ml import torch_export_torchscript
    exported_time = Timestamp(); exported_time.GetCurrentTime()
    spec=ModelSpec(
      source=Framework(package="torch", method="Framework_torch.torch_to_torchscript"),
      target=Framework(package="torch", method="Framework_torch.torchscript_to_torch"),
      folder=nbox_options.folder,
      name=nbox_options.model_name,
      requirements=["-f https://download.pytorch.org/whl/cpu/torch_stable.html","torch"], # contents inside requirements.txt file
      exported_time=exported_time,
      inputs=[],
      outputs=[],
    )
    spec=torch_export_torchscript(user_options, nbox_options, spec) # spec can be updated by the function
    return spec
  
  # @isthere("-f https://download.pytorch.org/whl/cpu/torch_stable.html","torch", soft=False)
  def torchscript_to_torch(user_options: TorchscriptToTorch, nbox_options: NboxOptions) -> Union[ModelSpec, Any]:
    '''
    Args:
      user_options (nbox.framework.ml.TorchscriptToTorch): Variables to be filled by the user
      _nbox_options (nbox.framework.NboxOptions): Options to be used by NBX-Services
    Returns:
      when serialising the model returns ``ModelSpec`` else returns the ``model`` in ``nbox.Model``
    '''
    from .on_ml import torch_import_torchscript
    exported_time = Timestamp(); exported_time.GetCurrentTime()
    spec=ModelSpec(
      source=Framework(package="torch", method="Framework_torch.torchscript_to_torch"),
      target=None, # this function has not been defined to throw any output
      folder=nbox_options.folder,
      name=nbox_options.model_name,
      requirements=["-f https://download.pytorch.org/whl/cpu/torch_stable.html","torch"], # contents inside requirements.txt file
      exported_time=exported_time,
      inputs=[],
      outputs=[],
    )
    spec=torch_import_torchscript(user_options, nbox_options, spec) # spec can be updated by the function
    return spec
  
  # @isthere("-f https://download.pytorch.org/whl/cpu/torch_stable.html","torch", soft=False)
  def torch_to_onnx(user_options: TorchToOnnx, nbox_options: NboxOptions) -> Union[ModelSpec, Any]:
    '''
    Args:
      user_options (nbox.framework.ml.TorchToOnnx): Variables to be filled by the user
      _nbox_options (nbox.framework.NboxOptions): Options to be used by NBX-Services
    Returns:
      when serialising the model returns ``ModelSpec`` else returns the ``model`` in ``nbox.Model``
    '''
    from .on_ml import torch_export_onnx
    exported_time = Timestamp(); exported_time.GetCurrentTime()
    spec=ModelSpec(
      source=Framework(package="torch", method="Framework_torch.torch_to_onnx"),
      target=Framework(package="onnx", method="Framework_onnx.onnx_runtime"),
      folder=nbox_options.folder,
      name=nbox_options.model_name,
      requirements=["-f https://download.pytorch.org/whl/cpu/torch_stable.html","torch"], # contents inside requirements.txt file
      exported_time=exported_time,
      inputs=[],
      outputs=[],
    )
    spec=torch_export_onnx(user_options, nbox_options, spec) # spec can be updated by the function
    return spec
  
  _METHODS={
    'torch_to_torchscript': (torch_to_torchscript, TorchToTorchscript),
    'torchscript_to_torch': (torchscript_to_torch, TorchscriptToTorch),
    'torch_to_onnx': (torch_to_onnx, TorchToOnnx),
    }


# End of file