Hello Suneeth. I would treat the partitioning behavior and the LazyInitializationException separately. For partitioning, userPartitions is the documented way to split the Entra user population using search filters. If performance is the goal, configure non-overlapping userPartitions and compare the aggregation runtime with your current setup.
For the exception, first reduce the Customization Rule to only: return object; Also check the Rule XML for any <ReferencedRules> entries.
Then run the same aggregation once with partitioning enabled and once with Enable Partitioning disabled. If the error occurs only with partitioning enabled, that would strongly point to an issue with the partitioned execution path rather than the code inside your rule.
In that case, collect the full stack trace and raise a SailPoint Support case. Since you are on 8.4p1, if you have a later supported 8.4 patch available in non-production, I would also try reproducing it there.
A similar Rule.referencedRules LazyInitializationException is worth comparing against.
Also, PIM is not unsupported with partitioning. SailPoint recommends non-partitioned aggregation when PIM is enabled because of the additional API calls. Since PIM is disabled in your environment, that does not apply here.