Index

A

Action (class in jaxarc) action (jaxarc.configs.JaxArcConfig attribute) (jaxarc.types.EnvParams attribute) action_space() (jaxarc.Environment method) (jaxarc.wrappers.BboxActionWrapper method) (jaxarc.wrappers.FlattenActionWrapper method) (jaxarc.wrappers.PointActionWrapper method) ActionConfig (class in jaxarc.configs)

allow_invalid_actions (jaxarc.configs.ActionConfig attribute) allowed_operations (jaxarc.configs.ActionConfig attribute) allowed_operations_mask (jaxarc.State attribute) AnswerObservationWrapper (class in jaxarc.wrappers) ARCOperationType (class in jaxarc.types) available_named_subsets() (in module jaxarc.registration) (jaxarc.registration.EnvRegistry method) available_task_ids() (in module jaxarc.registration) (jaxarc.registration.EnvRegistry method)

B

background_color (jaxarc.configs.DatasetConfig attribute) base_output_dir (jaxarc.configs.StorageConfig attribute) (jaxarc.utils.visualization.EpisodeConfig attribute)

batched_logging_enabled (jaxarc.configs.LoggingConfig attribute) BboxActionWrapper (class in jaxarc.wrappers) buffer (jaxarc.types.EnvParams attribute)

C

cleanup_old_data() (jaxarc.utils.visualization.EpisodeManager method) cleanup_policy (jaxarc.configs.StorageConfig attribute) (jaxarc.utils.visualization.EpisodeConfig attribute) CLEAR (jaxarc.types.ARCOperationType attribute) clear_output_on_start (jaxarc.configs.StorageConfig attribute) clipboard (jaxarc.State attribute) ClipboardObservationWrapper (class in jaxarc.wrappers) close() (jaxarc.Environment method) ColorValue (in module jaxarc.types) compress_old_episodes (jaxarc.utils.visualization.EpisodeConfig attribute) context_dependent_operations (jaxarc.configs.ActionConfig attribute)

ContextualObservationWrapper (class in jaxarc.wrappers) COPY (jaxarc.types.ARCOperationType attribute) COPY_INPUT (jaxarc.types.ARCOperationType attribute) create_episode_comparison_visualization() (in module jaxarc.utils.visualization) create_run_subdirs (jaxarc.configs.StorageConfig attribute) (jaxarc.utils.visualization.EpisodeConfig attribute) current_episode_dir (jaxarc.utils.visualization.EpisodeManager attribute) current_episode_num (jaxarc.utils.visualization.EpisodeManager attribute) current_run_dir (jaxarc.utils.visualization.EpisodeManager attribute) current_run_name (jaxarc.utils.visualization.EpisodeManager attribute) CUT (jaxarc.types.ARCOperationType attribute)

D

data (jaxarc.types.Grid attribute), [1] dataset (jaxarc.configs.JaxArcConfig attribute) (jaxarc.types.EnvParams attribute) dataset_name (jaxarc.configs.DatasetConfig attribute) dataset_path (jaxarc.configs.DatasetConfig attribute) dataset_repo (jaxarc.configs.DatasetConfig attribute) DatasetConfig (class in jaxarc.configs) debug_dir (jaxarc.configs.StorageConfig attribute) debug_level (jaxarc.configs.EnvironmentConfig attribute)

demo_weight (jaxarc.configs.GridInitializationConfig attribute) discount_space() (jaxarc.Environment method) DiscountValue (in module jaxarc.types) draw_episode_summary_svg() (in module jaxarc.utils.visualization) draw_grid_svg() (in module jaxarc.utils.visualization) draw_parsed_task_data_svg() (in module jaxarc.utils.visualization) draw_rl_step_svg() (in module jaxarc.utils.visualization) draw_task_pair_svg() (in module jaxarc.utils.visualization) dynamic_action_filtering (jaxarc.configs.ActionConfig attribute)

E

empty_weight (jaxarc.configs.GridInitializationConfig attribute) enabled (jaxarc.configs.VisualizationConfig attribute) (jaxarc.configs.WandbConfig attribute) entity (jaxarc.configs.WandbConfig attribute) env_entry (jaxarc.registration.EnvSpec attribute) Environment (class in jaxarc) environment (jaxarc.configs.JaxArcConfig attribute) EnvironmentConfig (class in jaxarc.configs) EnvParams (class in jaxarc.types)

EnvRegistry (class in jaxarc.registration) EnvSpec (class in jaxarc.registration) episode_dir_format (jaxarc.utils.visualization.EpisodeConfig attribute) episode_mode (jaxarc.types.EnvParams attribute) episode_summaries (jaxarc.configs.VisualizationConfig attribute) EpisodeConfig (class in jaxarc.utils.visualization) EpisodeManager (class in jaxarc.utils.visualization) episodes_dir (jaxarc.configs.StorageConfig attribute) estimate_storage_usage() (jaxarc.utils.visualization.EpisodeConfig method) expected_subdirs (jaxarc.configs.DatasetConfig attribute)

F

FILL_0 (jaxarc.types.ARCOperationType attribute) FILL_1 (jaxarc.types.ARCOperationType attribute) FILL_2 (jaxarc.types.ARCOperationType attribute) FILL_3 (jaxarc.types.ARCOperationType attribute) FILL_4 (jaxarc.types.ARCOperationType attribute) FILL_5 (jaxarc.types.ARCOperationType attribute) FILL_6 (jaxarc.types.ARCOperationType attribute) FILL_7 (jaxarc.types.ARCOperationType attribute) FILL_8 (jaxarc.types.ARCOperationType attribute) FILL_9 (jaxarc.types.ARCOperationType attribute) FlattenActionWrapper (class in jaxarc.wrappers) FLIP_HORIZONTAL (jaxarc.types.ARCOperationType attribute) FLIP_VERTICAL (jaxarc.types.ARCOperationType attribute) FLOOD_FILL_0 (jaxarc.types.ARCOperationType attribute) FLOOD_FILL_1 (jaxarc.types.ARCOperationType attribute) FLOOD_FILL_2 (jaxarc.types.ARCOperationType attribute) FLOOD_FILL_3 (jaxarc.types.ARCOperationType attribute) FLOOD_FILL_4 (jaxarc.types.ARCOperationType attribute)

FLOOD_FILL_5 (jaxarc.types.ARCOperationType attribute) FLOOD_FILL_6 (jaxarc.types.ARCOperationType attribute) FLOOD_FILL_7 (jaxarc.types.ARCOperationType attribute) FLOOD_FILL_8 (jaxarc.types.ARCOperationType attribute) FLOOD_FILL_9 (jaxarc.types.ARCOperationType attribute) force_cleanup_run() (jaxarc.utils.visualization.EpisodeManager method) from_config() (jaxarc.types.EnvParams class method) from_dict() (jaxarc.utils.visualization.EpisodeConfig class method) from_hydra() (jaxarc.configs.ActionConfig class method) (jaxarc.configs.DatasetConfig class method) (jaxarc.configs.EnvironmentConfig class method) (jaxarc.configs.GridInitializationConfig class method) (jaxarc.configs.JaxArcConfig class method) (jaxarc.configs.LoggingConfig class method) (jaxarc.configs.RewardConfig class method) (jaxarc.configs.StorageConfig class method) (jaxarc.configs.VisualizationConfig class method) (jaxarc.configs.WandbConfig class method) from_tuple() (jaxarc.utils.visualization.EpisodeConfig method)

G

generate_run_name() (jaxarc.utils.visualization.EpisodeConfig method) get_available_demo_pairs() (jaxarc.types.JaxArcTask method) get_available_test_pairs() (jaxarc.types.JaxArcTask method) get_base_path() (jaxarc.utils.visualization.EpisodeConfig method) get_config() (in module jaxarc.utils.core) get_current_run_info() (jaxarc.utils.visualization.EpisodeManager method) get_demo_pair_data() (jaxarc.types.JaxArcTask method) get_episode_summary_path() (jaxarc.utils.visualization.EpisodeManager method) get_grid_shape() (jaxarc.types.JaxArcTask method) get_max_test_pairs() (jaxarc.types.JaxArcTask method) get_max_train_pairs() (jaxarc.types.JaxArcTask method) get_path() (in module jaxarc.utils.core) get_raw_path() (in module jaxarc.utils.core) get_step_path() (jaxarc.utils.visualization.EpisodeManager method) get_subset_task_ids() (in module jaxarc.registration) (jaxarc.registration.EnvRegistry method)

get_task_id() (jaxarc.types.JaxArcTask method) get_task_summary() (jaxarc.types.JaxArcTask method) get_test_input_grid() (jaxarc.types.JaxArcTask method) get_test_output_grid() (jaxarc.types.JaxArcTask method) get_test_pair() (jaxarc.types.JaxArcTask method) get_test_pair_data() (jaxarc.types.JaxArcTask method) get_train_input_grid() (jaxarc.types.JaxArcTask method) get_train_output_grid() (jaxarc.types.JaxArcTask method) get_train_pair() (jaxarc.types.JaxArcTask method) Grid (class in jaxarc.types) grid_initialization (jaxarc.configs.JaxArcConfig attribute) (jaxarc.types.EnvParams attribute) GridArray (in module jaxarc.types) GridHeight (in module jaxarc.types) GridInitializationConfig (class in jaxarc.configs) GridWidth (in module jaxarc.types) group (jaxarc.configs.WandbConfig attribute)

I

id (jaxarc.registration.EnvSpec attribute) input_grid (jaxarc.State attribute) (jaxarc.types.TaskPair attribute), [1] input_grid_mask (jaxarc.State attribute) input_grids_examples (jaxarc.types.JaxArcTask attribute), [1]

input_masks_examples (jaxarc.types.JaxArcTask attribute), [1] InputGridObservationWrapper (class in jaxarc.wrappers) invalid_operation_policy (jaxarc.configs.ActionConfig attribute) is_demo_pair_available() (jaxarc.types.JaxArcTask method) is_test_pair_available() (jaxarc.types.JaxArcTask method) items() (jaxarc.utils.visualization.EpisodeConfig method)

J

jaxarc.configs module jaxarc.registration module jaxarc.types module

jaxarc.utils.core module jaxarc.utils.visualization module JaxArcConfig (class in jaxarc.configs) JaxArcTask (class in jaxarc.types) job_type (jaxarc.configs.WandbConfig attribute)

K

key (jaxarc.State attribute)

keys() (jaxarc.utils.visualization.EpisodeConfig method) kwargs (jaxarc.registration.EnvSpec attribute)

L

list_episodes_in_run() (jaxarc.utils.visualization.EpisodeManager method) load_all_subsets_for_dataset() (in module jaxarc.registration) load_from_file() (jaxarc.utils.visualization.EpisodeConfig class method) load_subset() (in module jaxarc.registration) load_subset_if_needed() (in module jaxarc.registration) log_format (jaxarc.configs.LoggingConfig attribute) log_frequency (jaxarc.configs.LoggingConfig attribute)

log_grid_to_console() (in module jaxarc.utils.visualization) log_level (jaxarc.configs.LoggingConfig attribute) log_operations (jaxarc.configs.LoggingConfig attribute) log_rewards (jaxarc.configs.LoggingConfig attribute) logging (jaxarc.configs.JaxArcConfig attribute) LoggingConfig (class in jaxarc.configs) logs_dir (jaxarc.configs.StorageConfig attribute)

M

make() (in module jaxarc) (in module jaxarc.registration) (jaxarc.registration.EnvRegistry method) mask (jaxarc.types.Grid attribute), [1] MaskArray (in module jaxarc.types) max_colors (jaxarc.configs.DatasetConfig attribute) max_episode_steps (jaxarc.configs.EnvironmentConfig attribute) (jaxarc.registration.EnvSpec attribute) (jaxarc.types.EnvParams attribute) max_episodes_per_run (jaxarc.configs.StorageConfig attribute) (jaxarc.utils.visualization.EpisodeConfig attribute) max_grid_height (jaxarc.configs.DatasetConfig attribute) max_grid_width (jaxarc.configs.DatasetConfig attribute) max_operations (jaxarc.configs.ActionConfig attribute) max_storage_gb (jaxarc.configs.StorageConfig attribute) (jaxarc.utils.visualization.EpisodeConfig attribute)

max_test_pairs (jaxarc.configs.DatasetConfig attribute) max_train_pairs (jaxarc.configs.DatasetConfig attribute) min_grid_height (jaxarc.configs.DatasetConfig attribute) min_grid_width (jaxarc.configs.DatasetConfig attribute) module jaxarc.configs jaxarc.registration jaxarc.types jaxarc.utils.core jaxarc.utils.visualization MOVE_DOWN (jaxarc.types.ARCOperationType attribute) MOVE_LEFT (jaxarc.types.ARCOperationType attribute) MOVE_RIGHT (jaxarc.types.ARCOperationType attribute) MOVE_UP (jaxarc.types.ARCOperationType attribute)

N

notes (jaxarc.configs.WandbConfig attribute)

num_test_pairs (jaxarc.types.JaxArcTask attribute), [1] num_train_pairs (jaxarc.types.JaxArcTask attribute), [1]

O

observation_shape() (jaxarc.Environment method) observation_space() (jaxarc.Environment method) ObservationArray (in module jaxarc.types) offline_mode (jaxarc.configs.WandbConfig attribute) operation (jaxarc.Action attribute), [1]

OperationId (in module jaxarc.types) OperationMask (in module jaxarc.types) output_grid (jaxarc.types.TaskPair attribute), [1] output_grids_examples (jaxarc.types.JaxArcTask attribute), [1] output_masks_examples (jaxarc.types.JaxArcTask attribute), [1]

P

pair_idx (jaxarc.State attribute) PairIndex (in module jaxarc.types) parser_entry_point (jaxarc.configs.DatasetConfig attribute) PASTE (jaxarc.types.ARCOperationType attribute) permutation_types (jaxarc.configs.GridInitializationConfig attribute)

permutation_weight (jaxarc.configs.GridInitializationConfig attribute) PointActionWrapper (class in jaxarc.wrappers) preserve_empty_dirs (jaxarc.utils.visualization.EpisodeConfig attribute) PRNGKey (in module jaxarc.types) project_name (jaxarc.configs.WandbConfig attribute)

R

random_density (jaxarc.configs.GridInitializationConfig attribute) random_pattern_type (jaxarc.configs.GridInitializationConfig attribute) random_weight (jaxarc.configs.GridInitializationConfig attribute) register() (in module jaxarc.registration) (jaxarc.registration.EnvRegistry method) register_subset() (in module jaxarc.registration) (jaxarc.registration.EnvRegistry method) render() (jaxarc.Environment method) (jaxarc.wrappers.StepVisualizationWrapper method) render_mode (jaxarc.configs.EnvironmentConfig attribute) replace() (jaxarc.utils.visualization.EpisodeConfig method)

reset() (jaxarc.Environment method) (jaxarc.wrappers.StepVisualizationWrapper method) RESIZE (jaxarc.types.ARCOperationType attribute) reward (jaxarc.configs.JaxArcConfig attribute) (jaxarc.types.EnvParams attribute) reward_space() (jaxarc.Environment method) RewardConfig (class in jaxarc.configs) RewardValue (in module jaxarc.types) ROTATE_C (jaxarc.types.ARCOperationType attribute) ROTATE_CC (jaxarc.types.ARCOperationType attribute) run_name (jaxarc.configs.StorageConfig attribute) (jaxarc.utils.visualization.EpisodeConfig attribute)

S

save_code (jaxarc.configs.WandbConfig attribute) save_rl_step_visualization() (in module jaxarc.utils.visualization) save_svg_drawing() (in module jaxarc.utils.visualization) save_to_file() (jaxarc.utils.visualization.EpisodeConfig method) selected (jaxarc.State attribute) selection (jaxarc.Action attribute), [1] selection_threshold (jaxarc.configs.ActionConfig attribute) SelectionArray (in module jaxarc.types) shape (jaxarc.types.Grid property) shuffle_tasks (jaxarc.configs.DatasetConfig attribute) similarity_score (jaxarc.State attribute) similarity_weight (jaxarc.configs.RewardConfig attribute) SimilarityScore (in module jaxarc.types) start_new_episode() (jaxarc.utils.visualization.EpisodeManager method) start_new_run() (jaxarc.utils.visualization.EpisodeManager method) State (class in jaxarc) state_space() (jaxarc.Environment method) step() (jaxarc.Environment method) (jaxarc.wrappers.BboxActionWrapper method) (jaxarc.wrappers.FlattenActionWrapper method) (jaxarc.wrappers.PointActionWrapper method) (jaxarc.wrappers.StepVisualizationWrapper method)

step_count (jaxarc.State attribute) step_file_format (jaxarc.utils.visualization.EpisodeConfig attribute) step_penalty (jaxarc.configs.RewardConfig attribute) step_visualizations (jaxarc.configs.VisualizationConfig attribute) StepCount (in module jaxarc.types) StepVisualizationWrapper (class in jaxarc.wrappers) storage (jaxarc.configs.JaxArcConfig attribute) StorageConfig (class in jaxarc.configs) structured_logging (jaxarc.configs.LoggingConfig attribute) SUBMIT (jaxarc.types.ARCOperationType attribute) subset_indices (jaxarc.types.EnvParams attribute) subset_task_ids() (in module jaxarc.registration) (jaxarc.registration.EnvRegistry method) success_bonus (jaxarc.configs.RewardConfig attribute)

T

tags (jaxarc.configs.WandbConfig attribute) target_grid (jaxarc.State attribute) target_grid_mask (jaxarc.State attribute) task_idx (jaxarc.State attribute) task_index (jaxarc.types.JaxArcTask attribute), [1] task_split (jaxarc.configs.DatasetConfig attribute) TaskIndex (in module jaxarc.types) TaskInputGrids (in module jaxarc.types) TaskInputMasks (in module jaxarc.types) TaskOutputGrids (in module jaxarc.types)

TaskOutputMasks (in module jaxarc.types) TaskPair (class in jaxarc.types) test_input_grids (jaxarc.types.JaxArcTask attribute), [1] test_input_masks (jaxarc.types.JaxArcTask attribute), [1] to_dict() (jaxarc.utils.visualization.EpisodeConfig method) to_tuple() (jaxarc.utils.visualization.EpisodeConfig method) to_yaml() (jaxarc.configs.JaxArcConfig method) to_yaml_file() (jaxarc.configs.JaxArcConfig method) true_test_output_grids (jaxarc.types.JaxArcTask attribute), [1] true_test_output_masks (jaxarc.types.JaxArcTask attribute), [1]

U

unsolved_submission_penalty (jaxarc.configs.RewardConfig attribute)

unwrapped (jaxarc.Environment property)

V

validate() (jaxarc.Action method) (jaxarc.configs.ActionConfig method) (jaxarc.configs.DatasetConfig method) (jaxarc.configs.EnvironmentConfig method) (jaxarc.configs.GridInitializationConfig method) (jaxarc.configs.JaxArcConfig method) (jaxarc.configs.LoggingConfig method) (jaxarc.configs.RewardConfig method) (jaxarc.configs.StorageConfig method) (jaxarc.configs.VisualizationConfig method) (jaxarc.configs.WandbConfig method)

validate_actions (jaxarc.configs.ActionConfig attribute) validate_storage_path() (jaxarc.utils.visualization.EpisodeConfig method) values() (jaxarc.utils.visualization.EpisodeConfig method) visualization (jaxarc.configs.JaxArcConfig attribute) visualization_dir (jaxarc.configs.StorageConfig attribute) VisualizationConfig (class in jaxarc.configs) visualize_grid_rich() (in module jaxarc.utils.visualization) visualize_parsed_task_data_rich() (in module jaxarc.utils.visualization) visualize_task_pair_rich() (in module jaxarc.utils.visualization)

W

wandb (jaxarc.configs.JaxArcConfig attribute) WandbConfig (class in jaxarc.configs)

working_grid (jaxarc.State attribute) working_grid_mask (jaxarc.State attribute)