Try adding Any

blackjax-devs · Jul 24, 2024 · 3149cf0 · 3149cf0
1 parent 8c5bb34
commit 3149cf0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/blackjax/base.py b/blackjax/base.py
@@ -10,7 +10,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Callable, NamedTuple, Optional
+from typing import Callable, NamedTuple, Optional, Any
 
 from typing_extensions import Protocol
 
@@ -64,7 +64,7 @@ class UpdateFn(Protocol):
 
     """
 
-    def __call__(self, rng_key: PRNGKey, state: State, *args) -> tuple[State, Info]:
+    def __call__(self, rng_key: PRNGKey, state: State, *args: Any) -> tuple[State, Info]:
         """Update the current state using the sampling algorithm.
 
         Parameters