Skip to content

Conversation

@kiranandcode
Copy link
Contributor

@kiranandcode kiranandcode commented Dec 9, 2025

Fixes #435 and closes #436.

    @effectful.ops.syntax.register_dataclass
    @dataclasses.dataclass
    class Point:
        x: int
        y: int

    @defop
    def random_point() -> Point:
        raise NotHandled

    def client():
        p1 = random_point()
        p2 = random_point()
        return p1.x + p2.x

    t = client()
    assert isinstance(t, Term)
    with handler({random_point: lambda: Point(1, 2)}):
        res = client()
        assert res == 2
        res = evaluate(t)
        assert res == 2

@kiranandcode
Copy link
Contributor Author

kiranandcode commented Dec 9, 2025

Let me rebase ontop of #438

actually #438 rebases onto staging-llm

Copy link
Contributor

@cooijmanstim cooijmanstim left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks a lot! We should also pass this by someone more familiar with the codebase such as @eb8680 or @jfeser.

@kiranandcode kiranandcode requested a review from jfeser December 10, 2025 17:25
Copy link
Contributor

@jfeser jfeser left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks generally good! A few nits noted.

@kiranandcode kiranandcode requested a review from jfeser December 10, 2025 20:28
@cooijmanstim
Copy link
Contributor

At this point, does it still need to be a metaclass? It seems like it could be a regular DataclassTerm class that is instantiated with a reference to the dataclass type and overrides __getattr__ to expose the properties.

@jfeser
Copy link
Contributor

jfeser commented Dec 10, 2025

Using a metaclass gives you isinstance(term, Dataclass).

@kiranandcode kiranandcode requested a review from eb8680 December 11, 2025 17:36
@kiranandcode
Copy link
Contributor Author

requesting additional eyes because this PR now touches a lot more of effectful's internals

@kiranandcode
Copy link
Contributor Author

didn't want to touch the core definitions so tried for a minimal as change as possible, but it might have been nicer to stuff the resolved type into Operation maybe?

@jfeser
Copy link
Contributor

jfeser commented Dec 11, 2025

It's not a property of the operation, but rather the operation applied to particular arguments. We shouldn't mutate the operation object to carry that data.

@kiranandcode
Copy link
Contributor Author

right, but if the operation applied to specific arguments changes its type (specializes it) then theres a reasonable interpretation for updating the operation to reflect the updated type.

that being said, I'm not familiar enough with the internals and usage of effectful to tell what that would break or implications of that, so if this current approach makes sense, then I have no objections.

@jfeser
Copy link
Contributor

jfeser commented Dec 11, 2025

So, we have a single operation object (e.g. identity) that can appear in many contexts in a term and have different return types. That object can't carry the type resulting from a single application.

@kiranandcode kiranandcode merged commit 4a9091d into master Dec 12, 2025
5 checks passed
@kiranandcode kiranandcode deleted the kg-dataclass-term branch December 12, 2025 02:05
@kiranandcode
Copy link
Contributor Author

Oh, sorry, was coding late and merged the branch but I realised I hadn't got @eb8680 's reviews on the code. Should I revert this?

@eb8680
Copy link
Contributor

eb8680 commented Dec 12, 2025

No, it's fine, @jfeser reviewed it

eb8680 added a commit that referenced this pull request Dec 12, 2025
* Release v0.2.3 (#374)

* Install prettyprinter for term when library is available (#386)

* install prettyprinter for term when library is available

* lint

* move code into types.py

* fix pypandoc issue (#397)

* Convert evaluate to a singledispatch (#398)

* convert evaluate to a singledispatch

* lint

* add jnp.pi and ArrayTerm.T (#394)

* Deprecate defterm (#399)

* deprecate defterm

* remove defterm case

* remove defterm

* lint

* evaluate distribution arguments

* lint

* remove interpreter

* Revert "remove interpreter"

This reverts commit 3044277.

* wip

* lint

* Rework numpyro distribution handling to enable symbolic distributions and handling of distribution methods (#311)

* refactor distribution operations

* add a test for typeof of distributions

* add tests for symbolic dists/arguments

* introduce operations for distribution methods

* comment

* fix tests

* work around #310

* replace hack with new hack

* tweak repr for _BaseOperation

* lint

* work around #312

* clean up access to dist ops

* wip

* wip

* add type annotations to get correct term conversion

* lint

* include distribution arguments as properties

* fix distribution calls

* try again

* fixes

* format

* Box the output of `__type_rule__` (#387)

* box the output of __type_rule__

* fix tests

* fix tests

* require callers of __type_rule__ to box arguments

* fix

* move Box out of ops.types

* lint

* fix test

* fix syntactic_eq implementation for jax arrays (#405)

* Fix recursion error in sizesof (#406)

* fix recursion error in sizesof

* format

* Allow `_BaseOperation` subclasses to have an overrideable `apply` method (#414)

* stash

* fixes

* initial

* wip

* lint

* ensure each subclass has a fresh operation

* wip

* wip

* lint

* wip

* wip

* lint

* refactor class method support

* move defops

* fix test

* remove singledispatch case and add test

* move definition

* cleanup

* simplify

* cleanup

* lint

* fix failing test

* fix classmethod

* __isabstractmethod__

* revert

---------

Co-authored-by: Eli <eli@basis.ai>

* Try pulling in pyproject.toml from staging-llm to master (#425)

* Generate instance-level `Operation`s for bound methods (#351)

* generalize __get__

* nits

* coverage of methoddescriptor api

* methodtype

* simplify

* simplify

* simplify

* format

* revert

* restore

* simplify

* simplify

* retain instance op on term construction

* Simplify apply inheritance

* assign

* put call next to init_subclass

* add explanatory comment

* Operation.apply -> Operation.__apply__

* add test based on issue description

* fix doctest

* Fix dataclass @defops and added dataclass metaclass (#439)

* fixed dataclass ordering and added metaclass for simplifying construction of dataclass terms

* ensure term fields are not being overriden

* added decorator and dataclass

* updated to make defdata registration automatic

* simplified dataclass loop

* updated to give property op an appropriate name

* added failing tests

* fixed failing test

* fixed numpyro/pyro/torch interfaces

* minor fix + test for deffn kwargs

---------

Co-authored-by: Jack Feser <jack.feser@gmail.com>
Co-authored-by: Tim Cooijmans <cooijmans.tim@gmail.com>
Co-authored-by: Kiran Gopinathan <23038502+kiranandcode@users.noreply.github.com>
eb8680 added a commit that referenced this pull request Dec 22, 2025
* Release v0.2.3 (#374)

* Install prettyprinter for term when library is available (#386)

* install prettyprinter for term when library is available

* lint

* move code into types.py

* fix pypandoc issue (#397)

* Convert evaluate to a singledispatch (#398)

* convert evaluate to a singledispatch

* lint

* add jnp.pi and ArrayTerm.T (#394)

* Deprecate defterm (#399)

* deprecate defterm

* remove defterm case

* remove defterm

* lint

* evaluate distribution arguments

* lint

* remove interpreter

* Revert "remove interpreter"

This reverts commit 3044277.

* wip

* lint

* Rework numpyro distribution handling to enable symbolic distributions and handling of distribution methods (#311)

* refactor distribution operations

* add a test for typeof of distributions

* add tests for symbolic dists/arguments

* introduce operations for distribution methods

* comment

* fix tests

* work around #310

* replace hack with new hack

* tweak repr for _BaseOperation

* lint

* work around #312

* clean up access to dist ops

* wip

* wip

* add type annotations to get correct term conversion

* lint

* include distribution arguments as properties

* fix distribution calls

* try again

* fixes

* format

* Box the output of `__type_rule__` (#387)

* box the output of __type_rule__

* fix tests

* fix tests

* require callers of __type_rule__ to box arguments

* fix

* move Box out of ops.types

* lint

* fix test

* fix syntactic_eq implementation for jax arrays (#405)

* Fix recursion error in sizesof (#406)

* fix recursion error in sizesof

* format

* Allow `_BaseOperation` subclasses to have an overrideable `apply` method (#414)

* stash

* fixes

* initial

* wip

* lint

* ensure each subclass has a fresh operation

* wip

* wip

* lint

* wip

* wip

* lint

* refactor class method support

* move defops

* fix test

* remove singledispatch case and add test

* move definition

* cleanup

* simplify

* cleanup

* lint

* fix failing test

* fix classmethod

* __isabstractmethod__

* revert

---------

Co-authored-by: Eli <eli@basis.ai>

* Try pulling in pyproject.toml from staging-llm to master (#425)

* Generate instance-level `Operation`s for bound methods (#351)

* generalize __get__

* nits

* coverage of methoddescriptor api

* methodtype

* simplify

* simplify

* simplify

* format

* revert

* restore

* simplify

* simplify

* retain instance op on term construction

* Simplify apply inheritance

* assign

* put call next to init_subclass

* add explanatory comment

* Operation.apply -> Operation.__apply__

* add test based on issue description

* fix doctest

* Fix dataclass @defops and added dataclass metaclass (#439)

* fixed dataclass ordering and added metaclass for simplifying construction of dataclass terms

* ensure term fields are not being overriden

* added decorator and dataclass

* updated to make defdata registration automatic

* simplified dataclass loop

* updated to give property op an appropriate name

* added failing tests

* fixed failing test

* fixed numpyro/pyro/torch interfaces

* minor fix + test for deffn kwargs

* Type check and lint example code (#449)

* format example code

* type check examples

* Add beam search example using thermometer continuations (#431)

* add beam search example using thermometer continuations

* address comments

* add docstring

* lint

* Fix for jax 0.8.2 (#455)

* fix for jax 0.8.2

* add more register

* format

---------

Co-authored-by: Jack Feser <jack.feser@gmail.com>
Co-authored-by: Tim Cooijmans <cooijmans.tim@gmail.com>
Co-authored-by: Kiran Gopinathan <23038502+kiranandcode@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Implement a general DataclassTerm for dataclass-valued terms Dataclasses broken

5 participants