content: # Copyright (c) Meta Platforms, Inc. and affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  import itertools
  import warnings
  from collections import OrderedDict
  from typing import (
      Any,
      Callable,
      Dict,
      Iterable,
      List,
      Optional,
      Protocol,
      Sequence,
      Tuple,
      Union,
  )
  
  import torch
  
  from theseus.constants import DeviceType
  from theseus.core.theseus_function import TheseusFunction
  from theseus.geometry.manifold import Manifold
  
  from .cost_function import CostFunction
  from .cost_weight import CostWeight
  from .variable import Variable
  
  
  class ErrorMetric(Protocol):
      def __call__(self, error_vector: torch.Tensor) -> torch.Tensor:
          pass
  
  
  def error_squared_norm_fn(error_vector: torch.Tensor) -> torch.Tensor:
      return (error_vector**2).sum(dim=1) / 2
  
  
  # If dtype is None, uses torch.get_default_dtype()
  class Objective:
      """An objective function to optimize.
  
      Defines the structure of an optimization problem in Theseus by aggregating
      :class:`cost functions <theseus.CostFunction>` into a single objective.
      The cost functions that comprise the final
      objective function are specified via the :meth:`add() <theseus.Objective.add>`
      method. Cost functions are responsible for registering their optimization and
      auxiliary variables, which are automatically added to the objective's list of
      variables when a cost function is added.
      Importantly, optimization variables must be instances of :class:`Manifold`
      subclasses, while auxiliary variables can be instances of
      any :class:`Variable` class.
  
      Args:
          dtype (optional[torch.dtype]): the data type to use for all variables. If
              ``None`` is passed, then uses ``torch.get_default_dtype()``.
          error_metric_fn (optional[callable]): a reference to a Python function used to
              aggregate cost functions into a single objective. Defaults to using the
              sum of squared costs. If given, it must receive a single tensor as input.
              The objective will use it to pass the batched concatenated error vector,
              will all cost function errors concatenated.
      """
  
      def __init__(
          self,
          dtype: Optional[torch.dtype] = None,
          error_metric_fn: Optional[ErrorMetric] = None,
          __allow_mixed_optim_aux_vars__: bool = False,  # experimental
      ):
          # maps variable names to the variable objects
          self.optim_vars: OrderedDict[str, Manifold] = OrderedDict()
  
          # maps variable names to variables objects, for optimization variables
          # that were registered when adding cost weights.
          self.cost_weight_optim_vars: OrderedDict[str, Manifold] = OrderedDict()
  
          # maps aux. variable names to the container objects
          self.aux_vars: OrderedDict[str, Variable] = OrderedDict()
  
          # maps variable name to variable, for any kind of variable added
          self._all_variables: OrderedDict[str, Variable] = OrderedDict()
  
          # maps cost function names to the cost function objects
          self.cost_functions: OrderedDict[str, CostFunction] = OrderedDict()
  
          # maps cost weights to the cost functions that use them
          # this is used when deleting cost function to check if the cost weight
          # variables can be deleted as well (when no other function uses them)
          self.cost_functions_for_weights: Dict[CostWeight, List[CostFunction]] = {}
  
          # ---- The following two methods are used just to get info from
          # ---- the objective, they don't affect the optimization logic.
          # a map from optimization variables to list of theseus functions it's
          # connected to
          self.functions_for_optim_vars: Dict[Manifold, List[TheseusFunction]] = {}
  
          # a map from all aux. variables to list of theseus functions it's connected to
          self.functions_for_aux_vars: Dict[Variable, List[TheseusFunction]] = {}
  
          self._batch_size: Optional[int] = None
  
          self.device: DeviceType = torch.device("cpu")
  
          self.dtype: Optional[torch.dtype] = dtype or torch.get_default_dtype()
  
          # this increases after every add/erase operation, and it's used to avoid
          # an optimizer to run on a stale version of the objective (since changing the
          # objective structure might break optimizer initialization).
          self.current_version = 0
  
          # ---- Callbacks for vectorization ---- #
          # This gets replaced when cost function vectorization is used.
          #
          # Normally, `_get_jacobians_iter()` returns an iterator over
          #  `self.cost_functions.values()`, so
          # that calling error() or jacobians() on the yielded cost functions
          # computes these quantities on demand.
          # But when vectorization is on, it will return this iterator that loops
          # over containers that serve cached jacobians and errors that have been
          # previously computed by the vectorization.
          self._vectorized_jacobians_iter: Optional[Iterable[CostFunction]] = None
  
          # Used to vectorize cost functions error + jacobians after an update
          # The results are cached so that the `self._get_jacobians_iter()` returns
          # them whenever called if no other updates have been done
          self._vectorization_run: Optional[Callable] = None
  
          # If vectorization is on, this gets replaced by a vectorized version
          # This method doesn't update the cache used by `self._get_jacobians_iter()`
          self._get_error_iter = self._get_error_iter_base
  
          # If vectorization is on, this will also handle vectorized containers
          self._vectorization_to: Optional[Callable] = None
  
          # If vectorization is on, this gets replaced by a vectorized version
          self._retract_method = Objective._retract_base
  
          # Keeps track of how many variable updates have been made to check
          # if vectorization should be updated
          self._num_updates_variables: Dict[str, int] = {}
  
          self._last_vectorization_has_grad = False
  
          self._vectorized = False
  
          # Computes an aggregation function for the error vector derived from costs
          # By default, this computes the squared norm of the error vector, divided by 2
          self._error_metric_fn = (
              error_metric_fn if error_metric_fn is not None else error_squared_norm_fn
          )
  
          self._allow_mixed_optim_aux_vars = __allow_mixed_optim_aux_vars__
  
          self.__tensor_ref = torch.empty(1, device=self.device, dtype=self.dtype)
  
      def _add_function_variables(
          self,
          function: TheseusFunction,
          optim_vars: bool = True,
          is_cost_weight: bool = False,
      ):
          if optim_vars:
              function_vars = function.optim_vars
              self_var_to_fn_map = self.functions_for_optim_vars
              self_vars_of_this_type = (
                  self.cost_weight_optim_vars if is_cost_weight else self.optim_vars
              )
          else:
              function_vars = function.aux_vars  # type: ignore
              self_var_to_fn_map = self.functions_for_aux_vars  # type: ignore
              self_vars_of_this_type = self.aux_vars  # type: ignore
  
          for variable in function_vars:
              # Check that variables have name and correct dtype
              if variable.name is None:
                  raise ValueError(
                      f"Variables added to an objective must be named, but "
                      f"{function.name} has an unnamed variable."
                  )
              if variable.dtype != self.dtype:
                  raise ValueError(
                      f"Tried to add variable {variable.name} with dtype "
                      f"{variable.dtype} but objective's dtype is {self.dtype}."
                  )
              # Check that names are unique
              if variable.name in self._all_variables:
                  if variable is not self._all_variables[variable.name]:
                      raise ValueError(
                          f"Two different variable objects with the "
                          f"same name ({variable.name}) are not allowed "
                          "in the same objective."
                      )
              else:
                  self._all_variables[variable.name] = variable
                  assert variable not in self_var_to_fn_map
                  self_var_to_fn_map[variable] = []
  
              # add to either self.optim_vars,
              # self.cost_weight_optim_vars or self.aux_vars
              self_vars_of_this_type[variable.name] = variable
  
              if self._allow_mixed_optim_aux_vars and variable not in self_var_to_fn_map:
                  self_var_to_fn_map[variable] = []
  
              # add to list of functions connected to this variable
              self_var_to_fn_map[variable].append(function)
  
      def add(self, cost_function: CostFunction):
          """Adds a cost function to the objective.
  
          When a cost function is added, this method goes over its list of registered
          optimization and auxiliary variables, and adds any of them to the objective's
          list of variables, as long as a variable with th4 same name hasn't been added
          before. If any of the cost function's variables has the same as that of
          a variable previously added to the objective, the method
          checks that they are referring to the same :class:`theseus.Variable`. If this
          is not the case, an error will be triggered. In other words, the objective
          expects to have a unique mapping between variable names and objects.
  
          The same procedure is followed for the cost function's weight.
  
          Args:
              cost_function (:class:`theseus.CostFunction`): the cost function to be
                  added to the objective.
  
          .. warning::
  
              If a cost weight registers optimization variables that are not used in any
              :class:`theseus.CostFunction <CostFunction>` objects, these will **NOT**
              be added to the set of the objective's optimization variables; they will be
              kept in a separate container. The :meth:`update` method will check for this,
              and throw a warning whenever this happens. Also note that Theseus
              always considers cost weights as constants, even if their value depends on
              variables declared as optimization variables.
          """
          # adds the cost function if not already present
          if cost_function.name in self.cost_functions:
              if cost_function is not self.cost_functions[cost_function.name]:
                  raise ValueError(
                      f"Two different cost function objects with the "
                      f"same name ({cost_function.name}) are not allowed "
                      "in the same objective."
                  )
              else:
                  warnings.warn(
                      "This cost function has already been added to the objective, "
                      "nothing to be done."
                  )
          else:
              self.cost_functions[cost_function.name] = cost_function
  
          self.current_version += 1
          # ----- Book-keeping for the cost function ------- #
          # adds information about the optimization variables in this cost function
          self._add_function_variables(cost_function, optim_vars=True)
  
          # adds information about the auxiliary variables in this cost function
          self._add_function_variables(cost_function, optim_vars=False)
  
          if cost_function.weight not in self.cost_functions_for_weights:
              # ----- Book-keeping for the cost weight ------- #
              # adds information about the variables in this cost function's weight
              self._add_function_variables(
                  cost_function.weight, optim_vars=True, is_cost_weight=True
              )
              # adds information about the auxiliary variables in this cost function's weight
              self._add_function_variables(
                  cost_function.weight, optim_vars=False, is_cost_weight=True
              )
  
              self.cost_functions_for_weights[cost_function.weight] = []
  
              if cost_function.weight.num_optim_vars() > 0:
                  raise RuntimeError(
                      f"The cost weight associated to {cost_function.name} receives one "
                      "or more optimization variables. Differentiating cost "
                      "weights with respect to optimization variables is not currently "
                      "supported, thus jacobians computed by our optimizers will be "
                      "incorrect. You may want to consider moving the weight computation "
                      "inside the cost function, so that the cost weight only receives "
                      "auxiliary variables."
                  )
  
          self.cost_functions_for_weights[cost_function.weight].append(cost_function)
  
          optim_vars_names = [
              var.name
              for var in itertools.chain(
                  cost_function.optim_vars, cost_function.weight.optim_vars
              )
          ]
          aux_vars_names = [
              var.name
              for var in itertools.chain(
                  cost_function.aux_vars, cost_function.weight.aux_vars
              )
          ]
          if not self._allow_mixed_optim_aux_vars:
              dual_var_err_msg = (
                  "Objective does not support a variable being both "
                  + "an optimization variable and an auxiliary variable."
              )
              for optim_name in optim_vars_names:
                  if self.has_aux_var(optim_name):
                      raise ValueError(dual_var_err_msg)
              for aux_name in aux_vars_names:
                  if self.has_optim_var(aux_name):
                      raise ValueError(dual_var_err_msg)
  
      def get_cost_function(self, name: str) -> CostFunction:
          """Returns a reference to the cost function with the given name.
  
          Args:
              name (str): the name of the cost function to retrieve.
  
          Returns:
              CostFunction: the :class:`theseus.CostFunction` with the given name, if
              present. Otherwise ``None``.
          """
          return self.cost_functions.get(name, None)
  
      def has_cost_function(self, name: str) -> bool:
          """Checks if a cost function with the given name is in the objective.
  
          Args:
              name (str): the name of the cost function.
  
          Returns:
              bool: ``True`` if a function exists with the given name,
                  ``False`` otherwise.
          """
          return name in self.cost_functions
  
      def has_optim_var(self, name: str) -> bool:
          """Checks if an optimization variable is used in the objective.
  
          Args:
              name (str): the name of the optimization variable.
  
          Returns:
              bool: ``True`` if an optimization variable with the given name exists in the
                  objective (i.e., it's currently associated to at least one cost
                  function in the objective). ``False`` otherwise.
          """
          return name in self.optim_vars
  
      def get_optim_var(self, name: str) -> Manifold:
          """Returns a reference to the optimization variable with the given name.
  
          Args:
              name (str): the name of the optimization variable to retrieve.
  
          Returns:
              Manifold: the :class:`theseus.Manifold` with the given name, if
              present. Otherwise ``None``.
          """
          return self.optim_vars.get(name, None)
  
      def has_aux_var(self, name: str) -> bool:
          """Checks if an auxiliary variable is used in the objective.
  
          Args:
              name (str): the name of the auxiliary variable.
  
          Returns:
              bool: ``True`` if an auxiliary variable with the given name exists in the
                  objective (i.e., it's currently associated to at least one cost
                  function or cost weight in the objective). ``False`` otherwise.
          """
          return name in self.aux_vars
  
      def get_aux_var(self, name: str) -> Variable:
          """Returns a reference to the auxiliary variable with the given name.
  
          Args:
              name (str): the name of the auxiliary variable to retrieve.
  
          Returns:
              Variable: the :class:`theseus.Variable` with the given name, if
              present. Otherwise ``None``.
          """
          return self.aux_vars.get(name, None)
  
      @property
      def batch_size(self) -> int:
          return self._batch_size
  
      def _erase_function_variables(
          self,
          function: TheseusFunction,
          optim_vars: bool = True,
          is_cost_weight: bool = False,
      ):
          if optim_vars:
              fn_var_list = function.optim_vars
              self_vars_of_this_type = (
                  self.cost_weight_optim_vars if is_cost_weight else self.optim_vars
              )
              self_var_to_fn_map = self.functions_for_optim_vars
          else:
              fn_var_list = function.aux_vars  # type: ignore
              self_vars_of_this_type = self.aux_vars  # type: ignore
              self_var_to_fn_map = self.functions_for_aux_vars  # type: ignore
  
          for variable in fn_var_list:
              cost_fn_idx = self_var_to_fn_map[variable].index(function)
              # remove function from the variable's list of connected cost functions
              del self_var_to_fn_map[variable][cost_fn_idx]
              # if the variable has no other functions, remove it also
              if not self_var_to_fn_map[variable]:
                  del self._all_variables[variable.name]
                  del self_var_to_fn_map[variable]
                  del self_vars_of_this_type[variable.name]
  
      def erase(self, name: str):
          """Removes a cost function from the objective given its name
  
          Also removes any of its variables that are no longer associated to other
          functions (either cost functions, or cost weights).
          Does the same for the cost weight variables, but only if the weight is
          not associated to any other cost function.
  
          Args:
              name (str): the name of the cost function to erase.
          """
          self.current_version += 1
          if name in self.cost_functions:
              cost_function = self.cost_functions[name]
              # erase variables associated to this cost function (if needed)
              self._erase_function_variables(cost_function, optim_vars=True)
              self._erase_function_variables(cost_function, optim_vars=False)
  
              # delete cost function from list of cost functions connected to its weight
              cost_weight = cost_function.weight
              cost_fn_idx = self.cost_functions_for_weights[cost_weight].index(
                  cost_function
              )
              del self.cost_functions_for_weights[cost_weight][cost_fn_idx]
  
              # No more cost functions associated to this weight, so can also delete
              if len(self.cost_functions_for_weights[cost_weight]) == 0:
                  # erase its variables (if needed)
                  self._erase_function_variables(
                      cost_weight, optim_vars=True, is_cost_weight=True
                  )
                  self._erase_function_variables(
                      cost_weight, optim_vars=False, is_cost_weight=True
                  )
                  del self.cost_functions_for_weights[cost_weight]
  
              # finally, delete the cost function
              del self.cost_functions[name]
          else:
              warnings.warn(
                  "This cost function is not in the objective, nothing to be done."
              )
  
      @staticmethod
      def _get_functions_connected_to_var(
          variable: Union[str, Variable],
          objectives_var_container_dict: "OrderedDict[str, Variable]",
          var_to_cost_fn_map: Dict[Variable, List[TheseusFunction]],
          variable_type: str,
      ) -> List[TheseusFunction]:
          if isinstance(variable, str):
              if variable not in objectives_var_container_dict:
                  raise ValueError(
                      f"{variable_type} named {variable} is not in the objective."
                  )
              variable = objectives_var_container_dict[variable]
          if variable not in var_to_cost_fn_map:
              raise ValueError(
                  f"{variable_type} {variable.name} is not in the objective."
              )
          return var_to_cost_fn_map[variable]
  
      def get_functions_connected_to_optim_var(
          self, variable: Union[Manifold, str]
      ) -> List[TheseusFunction]:
          """Gets a list of functions that depend on a given optimization variable.
  
          Args:
              variable (Union[Manifold, str]): the variable to query for. Can be an
                  instance of :class:`theseus.Manifold` or a string, specifying the name.
  
          Returns:
              List[TheseusFunction]: all cost functions that depend on the optimization
                  variable.
          """
          return Objective._get_functions_connected_to_var(
              variable,
              self.optim_vars,  # type: ignore
              self.functions_for_optim_vars,  # type: ignore
              "Optimization Variable",
          )
  
      def get_functions_connected_to_aux_var(
          self, aux_var: Union[Variable, str]
      ) -> List[TheseusFunction]:
          """Gets a list of functions that depend on a given auxiliary variable.
  
          Args:
              variable (Union[Variable, str]): the variable to query for. Can be an
                  instance of :class:`theseus.Variable` or a string, specifying the name.
  
          Returns:
              List[TheseusFunction]: all cost functions that depend on the auxiliary
                  variable.
          """
          return Objective._get_functions_connected_to_var(
              aux_var, self.aux_vars, self.functions_for_aux_vars, "Auxiliary Variable"
          )
  
      def dim(self) -> int:
          """Returns the dimension of the error vector.
  
          The dimension is equal to the sum of all cost functions' error dimensions.
  
          Returns:
              int: the error dimension.
          """
          err_dim = 0
          for cost_function in self.cost_functions.values():
              err_dim += cost_function.dim()
          return err_dim
  
      def size(self) -> Tuple[int, int]:
          """Returns the number of cost functions and variables in the objective.
  
          Returns:
              tuple[int, int]: the number of cost functions and optimization variables
                  in the objective, in that order.
          """
          return len(self.cost_functions), len(self.optim_vars)
  
      def size_cost_functions(self) -> int:
          """Returns the number of cost functions in the objective.
  
          Returns:
              int: the number of cost functions in the objective.
          """
          return len(self.cost_functions)
  
      def size_variables(self) -> int:
          """Returns the number of optimization variables in the objective.
  
          Returns:
              int: the number of optimization variables in the objective.
          """
          return len(self.optim_vars)
  
      def size_aux_vars(self) -> int:
          """Returns the number of auxiliary variables in the objective.
  
          Returns:
              int: the number of auxiliary variables in the objective.
          """
          return len(self.aux_vars)
  
      def error(
          self,
          input_tensors: Optional[Dict[str, torch.Tensor]] = None,
          also_update: bool = False,
      ) -> torch.Tensor:
          """Evaluates the error vector.
  
          Args:
              input_tensors (Dict[str, torch.Tensor], optional): if given, it must be a
                  dictionary mapping variable names to tensors; if a variable with the
                  given name is registered in the objective, its tensor will be replaced
                  with the one in the dictionary (possibly permanently, depending on the
                  value of ``also_update``). Defaults to ``None``, in which case the error
                  is evaluated using the current tensors stored in all registered
                  variables.
              also_update (bool, optional): if ``True``, and ``input_tensors`` is given,
                  the modified variables are permanently updated with the given tensors.
                  Defaults to ``False``, in which case the variables are reverted to the
                  previous tensors after the error is evaluated.
  
          Returns:
              torch.Tensor: a tensor of shape (batch_size x error_dim), with the
                  concatenation of all cost functions error vectors. The order corresponds
                  to the order in which cost functions were added to the objective.
          """
          old_tensors = {}
          if input_tensors is not None:
              if not also_update:
                  for var in self.optim_vars:
                      old_tensors[var] = self.optim_vars[var].tensor
              # Update vectorization only if the input tensors will be used for a
              # persistent update.
              self.update(input_tensors=input_tensors, _update_vectorization=also_update)
  
          # Current behavior when vectorization is on, is to always compute the error.
          # One could potentially optimize by only recompute when `input_tensors`` is
          # not None, and serving from the jacobians cache. However, when robust cost
          # functions are present this results in incorrect rescaling of error terms
          # so we are currently avoiding this optimization. Optimizers also compute error
          # by passing `input_tensors`, so for optimizers the current version should be
          # good enough.
          error_vector = torch.cat(
              [cf.weighted_error() for cf in self._get_error_iter()], dim=1
          )
  
          if input_tensors is not None and not also_update:
              # This line reverts back to the old tensors if a persistent update wasn't
              # required (i.e., `also_update is False`).
              # In this case, we pass _update_vectorization=False because
              # vectorization wasn't updated in the first call to `update()`.
              self.update(old_tensors, _update_vectorization=False)
          return error_vector
  
      def error_metric(
          self,
          input_tensors: Optional[Dict[str, torch.Tensor]] = None,
          also_update: bool = False,
      ) -> torch.Tensor:
          """Aggregates all cost function errors into a (batched) scalar objective.
  
          Args:
              input_tensors (Dict[str, torch.Tensor], optional): if given, it must be a
                  dictionary mapping variable names to tensors; if a variable with the
                  given name is registered in the objective, its tensor will be replaced
                  with the one in the dictionary (possibly permanently, depending on the
                  value of ``also_update``). Defaults to ``None``, in which case the error
                  is evaluated using the current tensors stored in all registered
                  variables.
              also_update (bool, optional): if ``True``, and ``input_tensors`` is given,
                  the modified variables are permanently updated with the given tensors.
                  Defaults to ``False``, in which case the variables are reverted to the
                  previous tensors after the error is evaluated.
  
          Returns:
              torch.Tensor: a tensor of shape (batch_size,) with the scalar value of
                  the objective function.
          """
          return self._error_metric_fn(
              self.error(input_tensors=input_tensors, also_update=also_update)
          )
  
      def copy(self) -> "Objective":
          """Creates a new copy of this objective.
  
          Returns:
              Objective: another instance of :class:`theseus.Objective` with copies
                  of all cost functions, weights, and variables, the same
                  connectivity structure, and error metric.
          """
          new_objective = Objective(
              dtype=self.dtype, error_metric_fn=self._error_metric_fn
          )
  
          # First copy all individual cost weights
          old_to_new_cost_weight_map: Dict[CostWeight, CostWeight] = {}
          for cost_weight in self.cost_functions_for_weights:
              new_cost_weight = cost_weight.copy(
                  new_name=cost_weight.name, keep_variable_names=True
              )
              old_to_new_cost_weight_map[cost_weight] = new_cost_weight
  
          # Now copy the cost functions and assign the corresponding cost weight copy
          new_cost_functions: List[CostFunction] = []
          for cost_function in self.cost_functions.values():
              new_cost_function = cost_function.copy(
                  new_name=cost_function.name, keep_variable_names=True
              )
              # we assign the allocated weight copies to avoid saving duplicates
              new_cost_function.weight = old_to_new_cost_weight_map[cost_function.weight]
              new_cost_functions.append(new_cost_function)
  
          # Handle case where a variable is copied in 2+ cost functions or cost weights,
          # since only a single copy should be maintained by objective
          for cost_function in new_cost_functions:
              # CostFunction
              for i, var in enumerate(cost_function.optim_vars):
                  if new_objective.has_optim_var(var.name):
                      cost_function.set_optim_var_at(
                          i, new_objective.optim_vars[var.name]
                      )
              for i, aux_var in enumerate(cost_function.aux_vars):
                  if new_objective.has_aux_var(aux_var.name):
                      cost_function.set_aux_var_at(
                          i, new_objective.aux_vars[aux_var.name]
                      )
              # CostWeight
              for i, var in enumerate(cost_function.weight.optim_vars):
                  if var.name in new_objective.cost_weight_optim_vars:
                      cost_function.weight.set_optim_var_at(
                          i, new_objective.cost_weight_optim_vars[var.name]
                      )
              for i, aux_var in enumerate(cost_function.weight.aux_vars):
                  if new_objective.has_aux_var(aux_var.name):
                      cost_function.weight.set_aux_var_at(
                          i, new_objective.aux_vars[aux_var.name]
                      )
              new_objective.add(cost_function)
          return new_objective
  
      def __deepcopy__(self, memo):
          if id(self) in memo:
              return memo[id(self)]
          the_copy = self.copy()
          memo[id(self)] = the_copy
          return the_copy
  
      def _resolve_batch_size(self):
          self._batch_size = None
  
          def _get_batch_size(batch_sizes: Sequence[int]) -> int:
              unique_batch_sizes = set(batch_sizes)
              if len(unique_batch_sizes) == 1:
                  return batch_sizes[0]
              if len(unique_batch_sizes) == 2:
                  min_bs = min(unique_batch_sizes)
                  max_bs = max(unique_batch_sizes)
                  if min_bs == 1:
                      return max_bs
              raise ValueError("Provided tensors must be broadcastable.")
  
          batch_sizes = [v.tensor.shape[0] for v in self.optim_vars.values()]
          batch_sizes.extend([v.tensor.shape[0] for v in self.aux_vars.values()])
          self._batch_size = _get_batch_size(batch_sizes)
  
      # batch_ignore_mask is a boolean list where batch_ignore_mask[i] = 1 means
      # for any variable v, v[i] will *not* be updated. Shape must be equal to the
      # batch size.
      def update(
          self,
          input_tensors: Optional[Dict[str, torch.Tensor]] = None,
          batch_ignore_mask: Optional[torch.Tensor] = None,
          _update_vectorization: bool = True,
      ):
          """Updates all variables with the given input tensor dictionary.
  
          The behavior of this method can be summarized by the following pseudocode:
  
          .. code-block::
  
              for name, tensor in input_tensors.items():
                  var = self.get_var_with_name(name).update(tensor)
              check_batch_size_consistency(self.all_variables)
  
          Any variables not included in the input tensors dictionary will retain their
          current tensors.
  
          After updating, the objective will modify its batch size
          property according to the resulting tensors. Therefore, all variable tensors
          must have a consistent batch size (either 1 or the same value as the others),
          after the update is completed. Note that this includes variables not referenced
          in the ``input_tensors`` dictionary.
  
          Args:
              input_tensors (Dict[str, torch.Tensor], optional): if given, it must be a
  