def _mlp(
        self,
        position: torch.Tensor,
        direction: Optional[torch.Tensor] = None,
        params: Optional[Dict[str, torch.Tensor]] = None,
        options: Optional[Dict[str, Any]] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        :param position: [batch_size x ... x 3]
        :param params: Meta parameters
        :param options: Optional hyperparameters
        :return: the final and directionless activations at the given query
        """
        params = self.update(params)
        options = AttrDict() if options is None else AttrDict(options)

        mlp = partial(self._run_mlp, direction=direction, params=params)
        parameters = []
        for i, layer in enumerate(self.mlp):
            if isinstance(layer, MetaLinear):
                parameters.extend(list(subdict(params, f"mlp.{i}").values()))
            else:
                parameters.extend(layer.parameters())

        h_final, h_directionless = checkpoint(
            mlp, (position,), parameters, options.checkpoint_stf_model
        )

        return h_final, h_directionless

        :param position: [batch_size x ... x 3]
        :param params: Meta parameters
        :param options: Optional hyperparameters
        :return: the final and directionless activations at the given query
        
shap_e/models/stf/mlp.py