Functions — github.com/brentyi/pips-jax

CodeHub MCPcopy Create free account

hub / github.com/brentyi/pips-jax / functions

Functions27 in github.com/brentyi/pips-jax

⨍ Functions 27 ◇ Types & classes 8

↓ 3 callersMethodinit_params

Returns network parameters.

src/pips_jax/model.py:44

↓ 2 callersMethod_feedforward

(self, x: Array, train: bool, name: str)

src/pips_jax/model.py:482

↓ 1 callersMethod_initialize_tracking

Initialize feature pyramids and particles for tracking.

src/pips_jax/model.py:132

↓ 1 callersMethod_iterative_update_step

Run a single iterative update step for our PIPs model.

src/pips_jax/model.py:170

↓ 1 callersMethod_make_layer

(self, planes: int, stride: int, name: str)

src/pips_jax/model.py:325

↓ 1 callersFunctionget_3d_embedding

Transformer-style positional encoding, in 3D.

src/pips_jax/model.py:496

↓ 1 callersFunctionrun_model

Run the PIPs model, and return coordinates at the final iteration of refinement.

Run full PIPs model. Returns a tuple of arrays: coords at each iter, visibility logits.

src/pips_jax/model.py:91

( self, rgbs: Float[Array, "b s h w 3"] )

src/pips_jax/model.py:249

( self, x: Float[Array, "... h w c"] )

src/pips_jax/model.py:347

( self, ffeats: Float[Array, "bn s latent_dim"], fcorrs: Float[Array, "bn s lww"],

src/pips_jax/model.py:403

( self, x: Float[Array, "bn s input_dim"], train: bool )

src/pips_jax/model.py:437

Method__enter__

benchmark.py:31

(self, *_unused_args)

benchmark.py:35

Functionbenchmark

( config: BenchmarkConfig, model: Pips = Pips(stride=4), )

benchmark.py:49

Functionbilerp_coords_batched_hw

Bilinear interpolation with leading batch axes.

src/pips_jax/utils_bilerp.py:13

Functionbilerp_coords_batched_hwc

Bilinear interpolation with leading batch axes and a trailing channel axis.

src/pips_jax/utils_bilerp.py:47

Functionbilerp_single

( features: Float[Array, "h w"], i: Float[Array, "*coord_batch"], j: Float[Array, "*co

src/pips_jax/utils_bilerp.py:24

( torch_checkpoint_path: Path = Path("./torch_reference_model/model-000200000.pth"), jax_checkpoint_pa

convert_checkpoint.py:17

( checkpoint_path: Path = Path("./checkpoints/reference_model/checkpoint_200000"), demo_image_dir: Pat

Functionmake_correlation_pyramids

Compute a correlation pyramid from a feature pyramid and target features.

src/pips_jax/utils_pyramid.py:32

Functionmake_feature_pyramids

Make a multi-resolution feature pyramid.

src/pips_jax/utils_pyramid.py:18

Methodrearrange_for_mixer

src/pips_jax/model.py:214

Functionresize_with_aligned_corners

Alternative to jax.image.resize(), which emulates align_corners=True in PyTorch's interpolation functions.

src/pips_jax/utils_bilerp.py:60

Functionsample_correlation_features

Sample a multi-resolution correlation feature from a correlation pyramid.

src/pips_jax/utils_pyramid.py:49

src/pips_jax/model.py:59

Methodupdate_step_wrapper

Wrapper for iterative refinement step, with a signature for use with `jax.lax.scan`.

src/pips_jax/model.py:108